fix(python): gemini batching, user agent and variable dims (#3618)

Carrying over from #2915, this patch introduces:
* Single-API call batching support for Gemini embeddings (up to 100 at a
time, the API limit)
* A versioned user agent header for Gemini API calls
* Support for [variable embedding dimension
size](https://ai.google.dev/gemini-api/docs/embeddings#control-embedding-size)
(Gemini is MRL trained)

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__}",
+                }
+            },
+        )

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