From 4d39a1a18fefebe140979f20e81f342a2ce96909 Mon Sep 17 00:00:00 2001
From: YutaSaito <36355491+uc4w6c@users.noreply.github.com>
Date: Sat, 6 Dec 2025 07:59:36 +0900
Subject: [PATCH] Fix: MLflow streaming spans for Anthropic passthrough
 (#17288)

* Fix: MLflow streaming spans for Anthropic passthrough

* fix: Revert "Handle MLflow chunk events without delta"
---
 litellm/integrations/mlflow.py                | 11 +++-
 .../test_litellm/integrations/test_mlflow.py  | 61 +++++++++++++++++++
 2 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py
index b348737868..6378e55f7e 100644
--- a/litellm/integrations/mlflow.py
+++ b/litellm/integrations/mlflow.py
@@ -129,8 +129,11 @@ class MlflowLogger(CustomLogger):
         self._add_chunk_events(span, response_obj)
 
         # If this is the final chunk, end the span. The final chunk
-        # has complete_streaming_response that gathers the full response.
-        if final_response := kwargs.get("complete_streaming_response"):
+        # has the assembled streaming response (key differs between sync/async paths).
+        final_response = kwargs.get("complete_streaming_response") or kwargs.get(
+            "async_complete_streaming_response"
+        )
+        if final_response:
             end_time_ns = int(end_time.timestamp() * 1e9)
 
             self._extract_and_set_chat_attributes(span, kwargs, final_response)
@@ -153,7 +156,9 @@ class MlflowLogger(CustomLogger):
                 span.add_event(
                     SpanEvent(
                         name="streaming_chunk",
-                        attributes={"delta": json.dumps(choice.delta.model_dump())},
+                        attributes={
+                            "delta": json.dumps(choice.delta.model_dump, default=str)
+                        },
                     )
                 )
         except Exception:
diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py
index b8894701e8..dba181def7 100644
--- a/tests/test_litellm/integrations/test_mlflow.py
+++ b/tests/test_litellm/integrations/test_mlflow.py
@@ -1,6 +1,8 @@
 import asyncio
+import json
 import os
 import sys
+from datetime import datetime
 from unittest.mock import MagicMock, patch
 
 # Adds the grandparent directory to sys.path to allow importing project modules
@@ -125,3 +127,62 @@ def test_mlflow_token_usage_attribute_structure():
             "output_tokens": 7,
             "total_tokens": 12,
         }
+
+
+def _mock_mlflow_modules():
+    mock_tracking = MagicMock()
+    mock_tracking.MlflowClient = MagicMock()
+
+    class DummySpanEvent:
+        def __init__(self, name, attributes):
+            self.name = name
+            self.attributes = attributes
+
+    mock_entities = MagicMock()
+    mock_entities.SpanStatusCode.OK = "OK"
+    mock_entities.SpanEvent = DummySpanEvent
+
+    return {
+        "mlflow": MagicMock(),
+        "mlflow.tracking": mock_tracking,
+        "mlflow.entities": mock_entities,
+        "mlflow.tracing.utils": MagicMock(),
+    }
+
+
+def test_mlflow_stream_handler_uses_async_complete_response():
+    modules = _mock_mlflow_modules()
+    with patch.dict("sys.modules", modules):
+        from litellm.integrations.mlflow import MlflowLogger
+
+        mlflow_logger = MlflowLogger()
+        mlflow_logger._start_span_or_trace = MagicMock(return_value="mock_span")
+        mlflow_logger._end_span_or_trace = MagicMock()
+        mlflow_logger._extract_and_set_chat_attributes = MagicMock()
+
+        class DummyDelta:
+            def model_dump(self, exclude_none=True):
+                return {"content": "chunk"}
+
+        response_obj = MagicMock()
+        response_obj.choices = [MagicMock(delta=DummyDelta())]
+
+        final_response = MagicMock()
+        kwargs = {
+            "litellm_call_id": "abc123",
+            "async_complete_streaming_response": final_response,
+        }
+
+        mlflow_logger._handle_stream_event(
+            kwargs=kwargs,
+            response_obj=response_obj,
+            start_time=datetime.utcnow(),
+            end_time=datetime.utcnow(),
+        )
+
+        mlflow_logger._end_span_or_trace.assert_called_once()
+        assert (
+            mlflow_logger._end_span_or_trace.call_args.kwargs["outputs"]
+            is final_response
+        )
+        assert "abc123" not in mlflow_logger._stream_id_to_span