# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors """Unit tests for GeminiText embedding function.""" import sys from unittest.mock import MagicMock, patch # Mock google.genai modules before they are imported by gemini_text.py mock_google = MagicMock() mock_genai = MagicMock() mock_types = MagicMock() mock_google.genai = mock_genai mock_genai.types = mock_types sys.modules["google"] = mock_google sys.modules["google.genai"] = mock_genai sys.modules["google.genai.types"] = mock_types import pytest # noqa: E402 import numpy as np # noqa: E402 from lancedb.embeddings import get_registry # noqa: E402 from lancedb import __version__ # noqa: E402 class TestGeminiText: """Tests for GeminiText model registration, configuration, and execution.""" @pytest.fixture(autouse=True) def setup_mocks(self): """Set up standard mocks for google-genai Client and Config.""" # Reset mocks mock_genai.reset_mock() mock_types.reset_mock() self.mock_client = MagicMock() mock_genai.Client.return_value = self.mock_client # Mock response for embed_content self.mock_embedding_1 = MagicMock() self.mock_embedding_1.values = [0.1] * 768 self.mock_embedding_2 = MagicMock() self.mock_embedding_2.values = [0.2] * 768 self.mock_response = MagicMock() self.mock_response.embeddings = [self.mock_embedding_1, self.mock_embedding_2] self.mock_client.models.embed_content.return_value = self.mock_response def test_gemini_registered(self): """Test that gemini-text is registered in the embedding function registry.""" registry = get_registry() assert registry.get("gemini-text") is not None def test_client_init_headers(self): """Test that Client is initialized with the partner-attribution header.""" with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}): with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"): registry = get_registry() func = registry.get("gemini-text").create() # Access the client property to trigger initialization _ = func.client mock_genai.Client.assert_called_once_with( api_key="test-key", http_options={ "headers": { "x-goog-api-client": f"lancedb/{__version__}", } }, ) def test_generate_embeddings_batched(self): """Test that multiple texts are sent in a single batched API request.""" with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}): with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"): registry = get_registry() func = registry.get("gemini-text").create() texts = ["hello", "world"] embeddings = func.generate_embeddings(texts) # Check embed_content was called exactly once self.mock_client.models.embed_content.assert_called_once() # Verify call arguments call_kwargs = self.mock_client.models.embed_content.call_args.kwargs assert call_kwargs["model"] == "gemini-embedding-001" assert len(call_kwargs["contents"]) == 2 assert call_kwargs["contents"][0] == {"parts": [{"text": "hello"}]} assert call_kwargs["contents"][1] == {"parts": [{"text": "world"}]} # Verify returns are correct numpy arrays assert len(embeddings) == 2 assert isinstance(embeddings[0], np.ndarray) assert embeddings[0].shape == (768,) assert np.allclose(embeddings[0], 0.1) assert np.allclose(embeddings[1], 0.2) def test_generate_embeddings_retrieval_document(self): """Test that retrieval_document task type prepends the document title part.""" with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}): with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"): registry = get_registry() func = registry.get("gemini-text").create( source_task_type="retrieval_document" ) texts = ["doc text"] # We need mock to return only 1 embedding since we only pass 1 text mock_embedding = MagicMock() mock_embedding.values = [0.3] * 768 self.mock_response.embeddings = [mock_embedding] embeddings = func.generate_embeddings( texts, task_type="retrieval_document" ) # Check call arguments for retrieval_document call_kwargs = self.mock_client.models.embed_content.call_args.kwargs assert call_kwargs["contents"][0] == { "parts": [{"text": "Embedding of a document"}, {"text": "doc text"}] } mock_types.EmbedContentConfig.assert_called_once_with( output_dimensionality=768, task_type="RETRIEVAL_DOCUMENT" ) assert len(embeddings) == 1 assert np.allclose(embeddings[0], 0.3) def test_custom_dimension(self): """Test that custom dimension (dim) can be configured and passed to config.""" with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}): with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"): registry = get_registry() func = registry.get("gemini-text").create(dim=3072) assert func.ndims() == 3072 texts = ["hello"] mock_embedding = MagicMock() mock_embedding.values = [0.5] * 3072 self.mock_response.embeddings = [mock_embedding] _ = func.generate_embeddings(texts) mock_types.EmbedContentConfig.assert_called_once_with( output_dimensionality=3072 ) def test_generate_embeddings_chunked(self): """Test that generate_embeddings chunks texts into groups of 100.""" with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}): with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"): registry = get_registry() func = registry.get("gemini-text").create() # Passing 250 texts should make 3 calls (100, 100, 50) texts = [f"text_{i}" for i in range(250)] # Mock client response to return correct number of embeddings per chunk def mock_embed_side_effect(model, contents, config=None): mock_resp = MagicMock() mock_embeddings = [] for _ in contents: emb = MagicMock() # Each embedding is length 768 emb.values = [0.1] * 768 mock_embeddings.append(emb) mock_resp.embeddings = mock_embeddings return mock_resp self.mock_client.models.embed_content.side_effect = ( mock_embed_side_effect ) embeddings = func.generate_embeddings(texts) # embed_content should be called 3 times assert self.mock_client.models.embed_content.call_count == 3 assert len(embeddings) == 250