From fc74cddb90950cbf82df9073780b0a4d495a1458 Mon Sep 17 00:00:00 2001
From: xiajiahang <xiajiahang82@gmail.com>
Date: Tue, 13 Oct 2026 20:13:29 +0800
Subject: [PATCH 1/2] fix(mlflow): classify embedding calls as embedding spans

---
 litellm/integrations/mlflow.py         |  2 +-
 tests/unit/integrations/test_mlflow.py | 25 +++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py
index 8731e96440f..44186db0b50 100644
--- a/litellm/integrations/mlflow.py
+++ b/litellm/integrations/mlflow.py
@@ -222,7 +222,7 @@ class MlflowLogger(CustomLogger):
 
         if call_type in ["completion", "acompletion"]:
             return SpanType.LLM
-        elif call_type == "embeddings":
+        elif call_type in ("embedding", "aembedding", "embeddings"):
             return SpanType.EMBEDDING
         else:
             return SpanType.LLM
diff --git a/tests/unit/integrations/test_mlflow.py b/tests/unit/integrations/test_mlflow.py
index f828c34a9ff..1d67e6a6a08 100644
--- a/tests/unit/integrations/test_mlflow.py
+++ b/tests/unit/integrations/test_mlflow.py
@@ -1,6 +1,7 @@
 import asyncio
 import json
 from datetime import datetime
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 # Adds the grandparent directory to sys.path to allow importing project modules
@@ -8,6 +9,30 @@ from unittest.mock import MagicMock, patch
 import pytest
 
 import litellm
+from litellm.types.utils import CallTypes
+
+
+@pytest.mark.parametrize(
+    ("call_type", "is_embedding"),
+    (
+        (CallTypes.embedding.value, True),
+        (CallTypes.aembedding.value, True),
+        ("embeddings", True),
+        (CallTypes.completion.value, False),
+        (CallTypes.acompletion.value, False),
+        (None, False),
+        ("unknown", False),
+    ),
+)
+def test_mlflow_span_type_matches_call_type(call_type: str | None, is_embedding: bool) -> None:
+    modules: Final = _mock_mlflow_modules()
+    with patch.dict("sys.modules", modules):
+        from litellm.integrations.mlflow import MlflowLogger
+
+        logger: Final = MlflowLogger()
+        span_types: Final = modules["mlflow.entities"].SpanType
+        expected: Final = span_types.EMBEDDING if is_embedding else span_types.LLM
+        assert logger._get_span_type(call_type) == expected
 
 
 @pytest.mark.asyncio

From fda2eb2a9be254af464cae9ed0f29578eb2aa91e Mon Sep 17 00:00:00 2001
From: xiajiahang <xiajiahang82@gmail.com>
Date: Tue, 13 Oct 2026 20:23:22 +0800
Subject: [PATCH 2/2] ci: rerun checks after unrelated MCP timeout