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 058d0ae885d..45bf36b32df 100644 --- a/tests/unit/integrations/test_mlflow.py +++ b/tests/unit/integrations/test_mlflow.py @@ -2,6 +2,7 @@ import asyncio import json import threading 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 @@ -9,6 +10,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