litellm/tests/unit/llms/volcengine/test_volcengine_embedding.py

"""Tests for Volcengine embedding."""

from unittest.mock import MagicMock, patch
import pytest

# Add parent directory to path for imports

import litellm
from litellm.types.utils import EmbeddingResponse


class TestVolcEngineEmbedding:
    """Test Volcengine embedding integration following LiteLLM patterns"""

    def get_custom_llm_provider(self) -> litellm.LlmProviders:
        return litellm.LlmProviders.VOLCENGINE

    def get_base_embedding_call_args(self) -> dict:
        return {
            "model": "volcengine/doubao-embedding-text-240715",
        }

    @pytest.mark.asyncio()
    @pytest.mark.parametrize("sync_mode", [True, False])
    async def test_basic_embedding(self, sync_mode):
        """Test basic embedding functionality with realistic response"""
        embedding_call_args = self.get_base_embedding_call_args()

        # Mock the embedding functions to avoid actual API calls
        with (
            patch("litellm.embedding") as mock_embedding,
            patch("litellm.aembedding") as mock_aembedding,
        ):
            # Create realistic Volcengine response
            mock_response = MagicMock()
            mock_response.model = "doubao-embedding-text-240715"
            mock_response.object = "list"
            mock_response.data = [
                {
                    "object": "embedding",
                    "embedding": [0.1, 0.2, 0.3]
                    + [0.01 * i for i in range(1021)],  # 1024-dim embedding
                    "index": 0,
                },
                {
                    "object": "embedding",
                    "embedding": [0.4, 0.5, 0.6]
                    + [0.02 * i for i in range(1021)],  # 1024-dim embedding
                    "index": 1,
                },
            ]
            mock_response.usage.prompt_tokens = 2
            mock_response.usage.total_tokens = 2

            mock_embedding.return_value = mock_response
            mock_aembedding.return_value = mock_response

            # Test sync mode
            if sync_mode is True:
                response = litellm.embedding(
                    **embedding_call_args,
                    input=["hello", "world"],
                )

                # Verify response structure matches Volcengine format
                assert response.model == "doubao-embedding-text-240715"
                assert response.object == "list"
                assert len(response.data) == 2
                assert len(response.data[0]["embedding"]) == 1024
                assert response.usage.total_tokens > 0

            # Test async mode
            else:
                response = await litellm.aembedding(
                    **embedding_call_args,
                    input=["hello", "world"],
                )

                # Verify response structure
                assert response.model == "doubao-embedding-text-240715"
                assert response.object == "list"
                assert len(response.data) == 2
                assert len(response.data[0]["embedding"]) == 1024
                assert response.usage.total_tokens > 0


def test_volcengine_embedding_with_encoding_formats():
    """Test Volcengine embedding with different encoding formats"""

    test_cases = [
        {"encoding_format": "float"},
        {"encoding_format": "base64"},
        {"encoding_format": None},  # Default
    ]

    for params in test_cases:
        with patch("litellm.embedding") as mock_embedding:
            # Create mock response based on encoding format
            mock_response = MagicMock()
            mock_response.model = "doubao-embedding-text-240715"
            mock_response.object = "list"

            if params["encoding_format"] == "base64":
                # Simulate base64 encoded embeddings
                mock_response.data = [
                    {
                        "object": "embedding",
                        "embedding": "c29tZS1iYXNlNjQtZW5jb2RlZC1lbWJlZGRpbmc=",  # base64 encoded
                        "index": 0,
                    }
                ]
            else:
                # Float embeddings (default)
                mock_response.data = [
                    {
                        "object": "embedding",
                        "embedding": [0.1, 0.2, 0.3, -0.1] * 256,  # 1024 dimensions
                        "index": 0,
                    }
                ]

            mock_response.usage.prompt_tokens = 3
            mock_response.usage.total_tokens = 3
            mock_embedding.return_value = mock_response

            # Test the call
            litellm.embedding(
                model="volcengine/doubao-embedding-text-240715",
                input=["test text"],
                **params,
            )

            # Verify the call was made with correct parameters
            mock_embedding.assert_called_once()
            call_args = mock_embedding.call_args
            assert call_args[1]["model"] == "volcengine/doubao-embedding-text-240715"
            assert call_args[1]["input"] == ["test text"]

            if params["encoding_format"] is not None:
                assert call_args[1]["encoding_format"] == params["encoding_format"]


def test_volcengine_embedding_with_user_parameter():
    """Test Volcengine embedding with user parameter for tracking"""

    with patch("litellm.embedding") as mock_embedding:
        mock_response = MagicMock()
        mock_response.model = "doubao-embedding-text-240715"
        mock_response.object = "list"
        mock_response.data = [
            {"object": "embedding", "embedding": [0.1] * 1024, "index": 0}
        ]
        mock_response.usage.prompt_tokens = 5
        mock_response.usage.total_tokens = 5
        mock_embedding.return_value = mock_response

        # Test with user parameter
        litellm.embedding(
            model="volcengine/doubao-embedding-text-240715",
            input=["user tracking test"],
            user="test-user-12345",
        )

        # Verify user parameter was passed
        mock_embedding.assert_called_once()
        call_args = mock_embedding.call_args
        assert call_args[1]["user"] == "test-user-12345"


def test_volcengine_embedding_error_scenarios():
    """Test Volcengine embedding error handling in integration context"""

    error_scenarios = [
        # Invalid model name
        {"model": "volcengine/invalid-model-name", "expected_error_pattern": "model"},
        # Invalid encoding format
        {
            "model": "volcengine/doubao-embedding-text-240715",
            "encoding_format": "invalid_format",
            "expected_error_pattern": "encoding_format",
        },
    ]

    for scenario in error_scenarios:
        with patch("litellm.embedding") as mock_embedding:
            # Configure mock to raise appropriate errors
            if "invalid-model" in scenario.get("model", ""):
                mock_embedding.side_effect = Exception("Model not found")
            elif scenario.get("encoding_format") == "invalid_format":
                mock_embedding.side_effect = ValueError("Unsupported encoding_format")

            # Test that errors are properly raised
            test_params = {
                k: v for k, v in scenario.items() if k != "expected_error_pattern"
            }

            with pytest.raises(Exception, match=f"(?i){scenario['expected_error_pattern']}") as exc_info:
                litellm.embedding(input=["test"], **test_params)

            # Verify error message contains expected pattern
            assert (
                scenario["expected_error_pattern"].lower()
                in str(exc_info.value).lower()
            )


def test_volcengine_embedding_with_multiple_inputs():
    """Test Volcengine embedding with various input lengths and types"""

    test_inputs = [
        # Single short text
        ["hello"],
        # Multiple short texts
        ["hello", "world", "test"],
        # Mixed length texts
        [
            "short",
            "This is a much longer text that should be handled properly by the embedding service",
        ],
        # Unicode content
        ["测试中文文本", "Test English text", "混合语言 mixed language"],
        # Many inputs (batch processing)
        [f"Test sentence number {i}" for i in range(10)],
    ]

    for test_input in test_inputs:
        with patch("litellm.embedding") as mock_embedding:
            # Create proportional mock response
            mock_response = MagicMock()
            mock_response.model = "doubao-embedding-text-240715"
            mock_response.object = "list"
            mock_response.data = [
                {
                    "object": "embedding",
                    "embedding": [0.1 * (i + 1)] * 1024,  # Unique embedding per input
                    "index": i,
                }
                for i in range(len(test_input))
            ]
            mock_response.usage.prompt_tokens = (
                len(test_input) * 5
            )  # Realistic token estimate
            mock_response.usage.total_tokens = len(test_input) * 5
            mock_embedding.return_value = mock_response

            # Test the call
            response = litellm.embedding(
                model="volcengine/doubao-embedding-text-240715", input=test_input
            )

            # Verify response matches input count
            assert len(response.data) == len(test_input)
            for i, embedding_data in enumerate(response.data):
                assert embedding_data["index"] == i
                assert len(embedding_data["embedding"]) == 1024


if __name__ == "__main__":
    pytest.main([__file__])