diff --git a/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py b/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py new file mode 100644 index 00000000000..c4afa7c6f12 --- /dev/null +++ b/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py @@ -0,0 +1,83 @@ +""" +Test OpenAI Embeddings Guardrail Translation Handler +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.llms.openai.embeddings.guardrail_translation.handler import ( + OpenAIEmbeddingsHandler, +) +from litellm.types.utils import CallTypes + + +@pytest.mark.asyncio +async def test_embeddings_handler_string_input(): + """Test embeddings handler with single string input""" + handler = OpenAIEmbeddingsHandler() + + # Mock guardrail + mock_guardrail = MagicMock() + mock_guardrail.apply_guardrail = AsyncMock(return_value={"texts": ["processed text"]}) + + data = { + "input": "Hello, world!", + "model": "text-embedding-3-small" + } + + result = await handler.process_input_messages( + data=data, + guardrail_to_apply=mock_guardrail, + ) + + # Verify guardrail was called with correct inputs + mock_guardrail.apply_guardrail.assert_called_once() + call_args = mock_guardrail.apply_guardrail.call_args + assert call_args.kwargs["inputs"]["texts"] == ["Hello, world!"] + assert call_args.kwargs["inputs"]["model"] == "text-embedding-3-small" + + # Verify result + assert result["input"] == "processed text" + + +@pytest.mark.asyncio +async def test_embeddings_handler_list_of_strings_input(): + """Test embeddings handler with list of strings input""" + handler = OpenAIEmbeddingsHandler() + + # Mock guardrail + mock_guardrail = MagicMock() + mock_guardrail.apply_guardrail = AsyncMock( + return_value={"texts": ["processed text 1", "processed text 2"]} + ) + + data = { + "input": ["Hello, world!", "How are you?"], + "model": "text-embedding-3-small" + } + + result = await handler.process_input_messages( + data=data, + guardrail_to_apply=mock_guardrail, + ) + + # Verify guardrail was called with correct inputs + mock_guardrail.apply_guardrail.assert_called_once() + call_args = mock_guardrail.apply_guardrail.call_args + assert call_args.kwargs["inputs"]["texts"] == ["Hello, world!", "How are you?"] + + # Verify result + assert result["input"] == ["processed text 1", "processed text 2"] + + +def test_embeddings_guardrail_translation_mappings(): + """Test that embeddings handler is registered for correct call types""" + from litellm.llms.openai.embeddings.guardrail_translation import ( + guardrail_translation_mappings, + ) + + assert CallTypes.embedding in guardrail_translation_mappings + assert CallTypes.aembedding in guardrail_translation_mappings + assert guardrail_translation_mappings[CallTypes.embedding] == OpenAIEmbeddingsHandler + assert guardrail_translation_mappings[CallTypes.aembedding] == OpenAIEmbeddingsHandler