diff --git a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/test_litellm/llms/vertex_ai/test_gemini_batch_embeddings.py similarity index 95% rename from tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py rename to tests/test_litellm/llms/vertex_ai/test_gemini_batch_embeddings.py index 54ea41a64504..98abf5459df7 100644 --- a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/tests/test_litellm/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -8,12 +8,8 @@ """ import json -import os -import sys from unittest.mock import MagicMock, patch -sys.path.insert(0, os.path.abspath("../../../..")) - import pytest import litellm @@ -311,13 +307,16 @@ def mock_auth_token(*args, **kwargs): ): mock_get_token.return_value = ( {"x-goog-api-key": "test-key"}, - "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent", + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:batchEmbedContents", ) mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = { - "embedding": {"values": [0.1, 0.2, 0.3, 0.4, 0.5]} + "embeddings": [ + {"values": [0.1, 0.2, 0.3, 0.4, 0.5]}, + {"values": [0.6, 0.7, 0.8, 0.9, 1.0]}, + ] } mock_post.return_value = mock_response @@ -338,17 +337,21 @@ def mock_auth_token(*args, **kwargs): request_body = json.loads(kwargs.get("data", "{}")) - assert "content" in request_body - assert "parts" in request_body["content"] - parts = request_body["content"]["parts"] + assert "requests" in request_body + assert len(request_body["requests"]) == 2 - assert len(parts) == 2 - assert parts[0]["text"] == "The food was delicious" - assert "inline_data" in parts[1] - assert parts[1]["inline_data"]["mime_type"] == "image/png" + text_parts = request_body["requests"][0]["content"]["parts"] + image_parts = request_body["requests"][1]["content"]["parts"] - assert len(response.data) == 1 + assert len(text_parts) == 1 + assert text_parts[0]["text"] == "The food was delicious" + assert len(image_parts) == 1 + assert "inline_data" in image_parts[0] + assert image_parts[0]["inline_data"]["mime_type"] == "image/png" + + assert len(response.data) == 2 assert response.data[0].embedding == [0.1, 0.2, 0.3, 0.4, 0.5] + assert response.data[1].embedding == [0.6, 0.7, 0.8, 0.9, 1.0] def test_gemini_multimodal_embedding_with_audio(): @@ -581,17 +584,21 @@ def mock_auth_token(*args, **kwargs): def test_filter_embed_params_drops_unsupported(): """Unsupported params like max_tokens should be filtered out.""" - result = _filter_embed_params({"dimensions": 768, "max_tokens": 256, "temperature": 0.5}) + result = _filter_embed_params( + {"dimensions": 768, "max_tokens": 256, "temperature": 0.5} + ) assert result == {"outputDimensionality": 768} def test_filter_embed_params_keeps_supported(): """All supported Gemini embedding params should pass through.""" - result = _filter_embed_params({ - "dimensions": 768, - "task_type": "RETRIEVAL_DOCUMENT", - "title": "My doc", - }) + result = _filter_embed_params( + { + "dimensions": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "My doc", + } + ) assert result == { "outputDimensionality": 768, "taskType": "RETRIEVAL_DOCUMENT", diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py similarity index 99% rename from tests/litellm/proxy/test_batch_x_litellm_model_encoding.py rename to tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index 49e0498f1409..101dc48603ad 100644 --- a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -422,9 +422,7 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_ model_id = "deployment-123" raw_batch_id = "batch_openai_123" - unified_batch_id = _make_unified_batch_id( - model_id=model_id, batch_id=raw_batch_id - ) + unified_batch_id = _make_unified_batch_id(model_id=model_id, batch_id=raw_batch_id) mock_response = _make_batch_response(batch_id=raw_batch_id, status="cancelled") mock_response._hidden_params = {} mock_router = MagicMock() diff --git a/tests/litellm/proxy/test_model_based_routing_files_batches.py b/tests/test_litellm/proxy/test_model_based_routing_files_batches.py similarity index 100% rename from tests/litellm/proxy/test_model_based_routing_files_batches.py rename to tests/test_litellm/proxy/test_model_based_routing_files_batches.py diff --git a/tests/litellm/test_batch_completion_models_all_responses.py b/tests/test_litellm/test_batch_completion_models_all_responses.py similarity index 100% rename from tests/litellm/test_batch_completion_models_all_responses.py rename to tests/test_litellm/test_batch_completion_models_all_responses.py