fix: fix wrong batching in Google GenAI Document Embedder (#2951)

* fix: fix wrong batching in Google GenAI Document Embedder

* unrelated: update supporte models

Author: anakin87

integrations/google_genai/src/haystack_integrations/components/embedders/google_genai/document_embedder.py

@@ -207,7 +207,7 @@ def _embed_batch(
         for batch in tqdm(
             batched(texts_to_embed, batch_size), disable=not self._progress_bar, desc="Calculating embeddings"
         ):
-            args: dict[str, Any] = {"model": self._model, "contents": [b[1] for b in batch]}
+            args: dict[str, Any] = {"model": self._model, "contents": list(batch)}
             if resolved_config:
                 args["config"] = resolved_config
 
@@ -238,7 +238,7 @@ async def _embed_batch_async(
         for batch in tqdm(
             batched(texts_to_embed, batch_size), disable=not self._progress_bar, desc="Calculating embeddings"
         ):
-            args: dict[str, Any] = {"model": self._model, "contents": [b[1] for b in batch]}
+            args: dict[str, Any] = {"model": self._model, "contents": list(batch)}
             if self._config:
                 args["config"] = types.EmbedContentConfig(**self._config) if self._config else None
 

integrations/google_genai/src/haystack_integrations/components/generators/google_genai/chat/chat_generator.py

@@ -153,6 +153,9 @@ def weather_function(city: str):
     """
 
     SUPPORTED_MODELS: ClassVar[list[str]] = [
+        "gemini-3.1-pro-preview",
+        "gemini-3-flash-preview",
+        "gemini-3.1-flash-lite-preview",
         "gemini-2.5-pro",
         "gemini-2.5-flash",
         "gemini-2.5-flash-lite",

integrations/google_genai/tests/test_document_embedder.py

@@ -4,12 +4,14 @@
 
 import os
 import random
+from unittest.mock import AsyncMock, MagicMock
 
+import numpy as np
 import pytest
 from haystack import Document
 from haystack.utils.auth import Secret
 
-from haystack_integrations.components.embedders.google_genai import GoogleGenAIDocumentEmbedder
+from haystack_integrations.components.embedders.google_genai import GoogleGenAIDocumentEmbedder, GoogleGenAITextEmbedder
 
 
 def mock_google_response(contents: list[str], model: str = "gemini-embedding-001", **kwargs) -> dict:
@@ -256,20 +258,66 @@ async def mock_embed_batch_async(texts_to_embed, batch_size):
         for doc_with_embedding in result["documents"]:
             assert doc_with_embedding.embedding == [0.1, 0.2, 0.3]
 
+    def test_embed_batch_passes_full_texts(self, monkeypatch):
+        monkeypatch.setenv("GOOGLE_API_KEY", "fake-api-key")
+        embedder = GoogleGenAIDocumentEmbedder(batch_size=2)
+
+        texts = ["first document text", "second document text", "third document text"]
+
+        mock_embedding = MagicMock()
+        mock_embedding.values = [0.1, 0.2, 0.3]
+
+        mock_response = MagicMock()
+        mock_response.embeddings = [mock_embedding]
+
+        embedder._client = MagicMock()
+        embedder._client.models.embed_content.return_value = mock_response
+
+        embedder._embed_batch(texts, batch_size=2)
+
+        calls = embedder._client.models.embed_content.call_args_list
+        assert len(calls) == 2
+        assert calls[0].kwargs["contents"] == ["first document text", "second document text"]
+        assert calls[1].kwargs["contents"] == ["third document text"]
+
+    @pytest.mark.asyncio
+    async def test_embed_batch_async_passes_full_texts(self, monkeypatch):
+        monkeypatch.setenv("GOOGLE_API_KEY", "fake-api-key")
+        embedder = GoogleGenAIDocumentEmbedder(batch_size=2)
+
+        texts = ["first document text", "second document text", "third document text"]
+
+        mock_embedding = MagicMock()
+        mock_embedding.values = [0.1, 0.2, 0.3]
+
+        mock_response = MagicMock()
+        mock_response.embeddings = [mock_embedding]
+
+        embedder._client = MagicMock()
+        embedder._client.aio.models.embed_content = AsyncMock(return_value=mock_response)
+
+        await embedder._embed_batch_async(texts, batch_size=2)
+
+        calls = embedder._client.aio.models.embed_content.call_args_list
+        assert len(calls) == 2
+        assert calls[0].kwargs["contents"] == ["first document text", "second document text"]
+        assert calls[1].kwargs["contents"] == ["third document text"]
+
     @pytest.mark.skipif(
         not os.environ.get("GOOGLE_API_KEY", None),
         reason="Export an env var called GOOGLE_API_KEY containing the Google API key to run this test.",
     )
     @pytest.mark.integration
     def test_run(self):
+        model = "gemini-embedding-001"
+
         docs = [
-            Document(content="I love cheese", meta={"topic": "Cuisine"}),
-            Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}),
+            Document(content="The capybara is the largest rodent in the world and lives near rivers in South America."),
+            Document(content="Dogs are domesticated mammals known for their loyalty and bond with humans."),
+            Document(content="The tiger is the largest big cat, recognized by its orange coat with black stripes."),
         ]
 
-        model = "gemini-embedding-001"
-
-        embedder = GoogleGenAIDocumentEmbedder(model=model, meta_fields_to_embed=["topic"], embedding_separator=" | ")
+        embedder = GoogleGenAIDocumentEmbedder(model=model, config={"task_type": "RETRIEVAL_DOCUMENT"})
 
         result = embedder.run(documents=docs)
         documents_with_embeddings = result["documents"]
@@ -283,6 +331,18 @@ def test_run(self):
 
         assert result["meta"]["model"] == model
 
+        text_embedder = GoogleGenAITextEmbedder(model=model, config={"task_type": "RETRIEVAL_QUERY"})
+        query_embedding = text_embedder.run("capybara")["embedding"]
+        query_vec = np.array(query_embedding)
+
+        similarities = []
+        for doc in documents_with_embeddings:
+            doc_vec = np.array(doc.embedding)
+            cosine_sim = np.dot(query_vec, doc_vec) / (np.linalg.norm(query_vec) * np.linalg.norm(doc_vec))
+            similarities.append(cosine_sim)
+
+        assert similarities[0] == max(similarities)
+
     @pytest.mark.asyncio
     @pytest.mark.skipif(
         not os.environ.get("GOOGLE_API_KEY", None),