diff --git a/python/python/lancedb/embeddings/gemini_text.py b/python/python/lancedb/embeddings/gemini_text.py index a9b10b002..1102e92d1 100644 --- a/python/python/lancedb/embeddings/gemini_text.py +++ b/python/python/lancedb/embeddings/gemini_text.py @@ -4,7 +4,7 @@ import os from functools import cached_property -from typing import List, Union +from typing import List, Optional, Union import numpy as np @@ -15,6 +15,8 @@ from .base import TextEmbeddingFunction from .registry import register from .utils import TEXT, api_key_not_found_help +EMBEDDING_BATCH_SIZE = 100 + @register("gemini-text") class GeminiText(TextEmbeddingFunction): @@ -81,6 +83,7 @@ class GeminiText(TextEmbeddingFunction): """ name: str = "gemini-embedding-001" + dim: Optional[int] = None query_task_type: str = "retrieval_query" source_task_type: str = "retrieval_document" @@ -93,6 +96,8 @@ class GeminiText(TextEmbeddingFunction): model_config["ignored_types"] = (cached_property,) def ndims(self): + if self.dim: + return self.dim # TODO: fix hardcoding return 768 @@ -133,22 +138,22 @@ class GeminiText(TextEmbeddingFunction): contents.append({"parts": [{"text": text}]}) # Build config - config_kwargs = {} + config_kwargs = {"output_dimensionality": self.ndims()} if task_type: config_kwargs["task_type"] = task_type.upper() # API expects uppercase - # Call embed_content for each content + config = types.EmbedContentConfig(**config_kwargs) if config_kwargs else None + + # Call embed_content in groups of at most EMBEDDING_BATCH_SIZE docs at a time embeddings = [] - for content in contents: - config = ( - types.EmbedContentConfig(**config_kwargs) if config_kwargs else None - ) + for i in range(0, len(contents), EMBEDDING_BATCH_SIZE): + chunk = contents[i : i + EMBEDDING_BATCH_SIZE] response = self.client.models.embed_content( model=self.name, - contents=content, + contents=chunk, config=config, ) - embeddings.append(response.embeddings[0].values) + embeddings.extend([np.array(e.values) for e in response.embeddings]) return embeddings @@ -160,5 +165,13 @@ class GeminiText(TextEmbeddingFunction): api_key_not_found_help("google") from google import genai as genai_module + from lancedb import __version__ - return genai_module.Client(api_key=os.environ.get("GOOGLE_API_KEY")) + return genai_module.Client( + api_key=os.environ.get("GOOGLE_API_KEY"), + http_options={ + "headers": { + "x-goog-api-client": f"lancedb/{__version__}", + } + }, + ) diff --git a/python/python/tests/test_gemini.py b/python/python/tests/test_gemini.py new file mode 100644 index 000000000..4e803095e --- /dev/null +++ b/python/python/tests/test_gemini.py @@ -0,0 +1,183 @@ +# 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