From a173232f8c2ed2fee9a33acee07424f5cb7553de Mon Sep 17 00:00:00 2001
From: Sam Chou <sam.chou@windystudios.com>
Date: Fri, 24 Oct 2025 12:57:15 -0700
Subject: [PATCH] Fix MLFlow tags - split request_tags into (key, val) if
 request_tag has colon (#15914)

* Fix mlflow tags - split request_tags into (key, val) if request_tag has colon

* Redundant name: tag_dict -> tags
---
 litellm/integrations/mlflow.py                | 26 ++++++++++++++++---
 .../test_litellm/integrations/test_mlflow.py  | 10 +++++--
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py
index 86af800d73..b348737868 100644
--- a/litellm/integrations/mlflow.py
+++ b/litellm/integrations/mlflow.py
@@ -60,7 +60,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"):
@@ -184,7 +187,9 @@ class MlflowLogger(CustomLogger):
             "call_type": kwargs.get("call_type"),
             "model": kwargs.get("model"),
         }
-        standard_obj: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object")
+        standard_obj: Optional[StandardLoggingPayload] = kwargs.get(
+            "standard_logging_object"
+        )
         if standard_obj:
             attributes.update(
                 {
@@ -257,12 +262,25 @@ 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,
             )
 
     def _transform_tag_list_to_dict(self, tag_list: list) -> dict:
-        return {tag: "" for tag in tag_list}
+        """
+        Transform a list of colon-separated tags into a dictionary.
+        Tags without colons are stored with empty string as the value.
+        """
+        tags = {}
+        for tag in tag_list:
+            if ":" in tag:
+                k, v = tag.split(":", 1)
+                tags[k.strip()] = v.strip()
+            else:
+                tags[tag.strip()] = ""
+        return tags
 
     def _end_span_or_trace(self, span, outputs, end_time_ns, status):
         """End an MLflow span or a trace."""
diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py
index 3b75268d55..b8894701e8 100644
--- a/tests/test_litellm/integrations/test_mlflow.py
+++ b/tests/test_litellm/integrations/test_mlflow.py
@@ -57,7 +57,7 @@ async def test_mlflow_logging_functionality():
             messages=[{"role": "user", "content": "test message"}],
             prediction=test_prediction,
             mock_response="test response",
-            metadata={"tags": ["tag1", "tag2", "production"]},
+            metadata={"tags": ["tag1", "tag2", "production", "jobID:214590dsff09fds", "taskName:run_page_classification"]},
         )
 
         # Allow time for async processing
@@ -72,7 +72,13 @@ async def test_mlflow_logging_functionality():
 
         # Check that tags parameter was included and properly transformed
         tags_param = call_args.kwargs.get("tags", {})
-        expected_tags = {"tag1": "", "tag2": "", "production": ""}
+        expected_tags = {
+            "tag1": "",
+            "tag2": "",
+            "production": "",
+            "jobID": "214590dsff09fds",
+            "taskName": "run_page_classification",
+        }
         assert tags_param == expected_tags, f"Expected tags {expected_tags}, got {tags_param}"
 
         # Check that prediction parameter was included in inputs