tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 10:21:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix MLFlow tags - split request_tags into (key, val) if request_tag has colon (#15914)

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

* Redundant name: tag_dict -> tags
This commit is contained in:
Sam Chou
2025-10-24 12:57:15 -07:00
committed by GitHub
parent 8b1424166b
commit a173232f8c
2 changed files with 30 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/mlflow.py
+22 -4
View File
@@ -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."""
tests/test_litellm/integrations/test_mlflow.py
+8 -2
View File
@@ -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