Include predicted output in tracing

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

litellm/integrations/mlflow.py

@@ -60,10 +60,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"):
@@ -168,6 +165,11 @@ class MlflowLogger(CustomLogger):
         for key in ["functions", "tools", "stream", "tool_choice", "user"]:
             if value := kwargs.get("optional_params", {}).pop(key, None):
                 inputs[key] = value
+
+        # Include prediction parameter if present
+        if prediction := kwargs.get("prediction"):
+            inputs["prediction"] = prediction
+
         return inputs
 
     def _extract_attributes(self, kwargs):
@@ -232,7 +234,6 @@ 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)
@@ -260,6 +261,7 @@ class MlflowLogger(CustomLogger):
                 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}
 

tests/test_litellm/integrations/test_mlflow.py

@@ -12,66 +12,75 @@ import litellm
 
 
 @pytest.mark.asyncio
-async def test_mlflow_request_tags_functionality():
-    """Test that request_tags are properly extracted and transformed into tags for MLflow traces."""
-    
+async def test_mlflow_logging_functionality():
+    """Test that request_tags and prediction parameters are properly logged in 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
-    
+
     # Mock all MLflow-related imports to avoid requiring MLflow as a dependency
     mock_mlflow_tracking = MagicMock()
     mock_mlflow_tracking.MlflowClient = MagicMock(return_value=mock_client)
-    
+
     mock_mlflow_entities = MagicMock()
     mock_mlflow_entities.SpanStatusCode.OK = "OK"
     mock_mlflow_entities.SpanStatusCode.ERROR = "ERROR"
     mock_mlflow_entities.SpanType.LLM = "LLM"
-    
+
     mock_mlflow = MagicMock()
     mock_mlflow.get_current_active_span.return_value = None
-    
-    with patch.dict('sys.modules', {
-        'mlflow': mock_mlflow,
-        'mlflow.tracking': mock_mlflow_tracking,
-        'mlflow.entities': mock_mlflow_entities,
-        'mlflow.tracing.utils': MagicMock(),
-    }):
+
+    with patch.dict(
+        "sys.modules",
+        {
+            "mlflow": mock_mlflow,
+            "mlflow.tracking": mock_mlflow_tracking,
+            "mlflow.entities": mock_mlflow_entities,
+            "mlflow.tracing.utils": MagicMock(),
+        },
+    ):
         # Now we can safely import MlflowLogger
         from litellm.integrations.mlflow import MlflowLogger
 
         # Create MlflowLogger instance
         mlflow_logger = MlflowLogger()
         litellm.callbacks = [mlflow_logger]
-        
-        # Test completion with request_tags
+
+        # Test completion with request_tags and prediction parameter
+        test_prediction = {"type": "content", "content": "This is a predicted output"}
         await litellm.acompletion(
             model="gpt-3.5-turbo",
             messages=[{"role": "user", "content": "test message"}],
+            prediction=test_prediction,
             mock_response="test response",
-            metadata={
-                "tags": ["tag1", "tag2", "production"]
-            }
+            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', {})
+        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}"
 
+        # 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']}"
+        )
 
 
 def test_mlflow_token_usage_attribute_structure():