chore(genai): Update google-genai dependency to support 2.x (#1766)

Author: sararob

libs/genai/langchain_google_genai/chat_models.py

@@ -394,10 +394,16 @@ def _convert_to_parts(
                             part_kwargs["media_resolution"] = {
                                 "level": part["media_resolution"]
                             }
+                    thought_signature = None
                     if "extras" in part and isinstance(part["extras"], dict):
                         sig = part["extras"].get("signature")
-                        if sig and isinstance(sig, str):
-                            part_kwargs["thought_signature"] = base64.b64decode(sig)
+                        if isinstance(sig, str):
+                            thought_signature = base64.b64decode(sig)
+                        elif isinstance(sig, bytes):
+                            thought_signature = sig
+
+                        if thought_signature:
+                            pass
 
                     parts.append(Part(**part_kwargs))
                 elif part["type"] == "image_url":
@@ -411,11 +417,13 @@ def _convert_to_parts(
                     # Check for thought_signature in extras
                     # (needed for multi-turn image editing/usage)
                     thought_sig = None
+                    image_part = image_loader.load_part(img_url)
                     if "extras" in part and isinstance(part["extras"], dict):
                         sig = part["extras"].get("signature")
-                        if sig and isinstance(sig, str):
-                            thought_sig = base64.b64decode(sig)
-                    image_part = image_loader.load_part(img_url)
+                        if isinstance(sig, str):
+                            image_part.thought_signature = base64.b64decode(sig)
+                        elif isinstance(sig, bytes):
+                            image_part.thought_signature = sig
                     if thought_sig:
                         image_part.thought_signature = thought_sig
                     parts.append(image_part)
@@ -429,12 +437,23 @@ def _convert_to_parts(
                     media_part_kwargs: dict[str, Any] = {}
 
                     if "data" in part:
-                        # Embedded media
+                        data = part["data"]
+                        if isinstance(data, str):
+                            clean_data = re.sub(r"\s+", "", data)
+                            data_validation_msg = "Data should be valid base64"
+                            if (
+                                not re.match(r"^[A-Za-z0-9+/]*={0,2}$", clean_data)
+                                or len(clean_data) % 4 != 0
+                            ):
+                                raise ValueError(data_validation_msg)
+                            try:
+                                data = base64.b64decode(clean_data)
+                            except Exception:
+                                raise ValueError(data_validation_msg)
                         media_part_kwargs["inline_data"] = Blob(
-                            data=part["data"], mime_type=mime_type
+                            data=data, mime_type=mime_type
                         )
                     elif "file_uri" in part:
-                        # Referenced files (e.g. stored in GCS)
                         media_part_kwargs["file_data"] = FileData(
                             file_uri=part["file_uri"], mime_type=mime_type
                         )
@@ -445,6 +464,12 @@ def _convert_to_parts(
                         _validate_video_metadata(part["video_metadata"])
                         metadata = VideoMetadata.model_validate(part["video_metadata"])
                         media_part_kwargs["video_metadata"] = metadata
+                        if "extras" in part and isinstance(part["extras"], dict):
+                            sig = part["extras"].get("signature")
+                            if sig and isinstance(sig, str):
+                                media_part_kwargs["thought_signature"] = (
+                                    base64.b64decode(sig)
+                                )
 
                     if "media_resolution" in part:
                         if model and _is_gemini_25_model(model):
@@ -461,10 +486,12 @@ def _convert_to_parts(
                             }
                     if "extras" in part and isinstance(part["extras"], dict):
                         sig = part["extras"].get("signature")
-                        if sig and isinstance(sig, str):
+                        if isinstance(sig, str):
                             media_part_kwargs["thought_signature"] = base64.b64decode(
                                 sig
                             )
+                        elif isinstance(sig, bytes):
+                            media_part_kwargs["thought_signature"] = sig
 
                     parts.append(Part(**media_part_kwargs))
                 elif part["type"] == "thinking":
@@ -993,7 +1020,7 @@ def _parse_response_candidate(
     effective_model_name = model_name_for_content or model_name
 
     parts = response_candidate.content.parts or [] if response_candidate.content else []
-    for part in parts:
+    for i, part in enumerate(parts):
         text: str | None = None
         try:
             if hasattr(part, "text") and part.text is not None:
@@ -1122,7 +1149,8 @@ def _parse_response_candidate(
             )
             additional_kwargs["function_call"] = function_call
 
-            tool_call_id = function_call.get("id", str(uuid.uuid4()))
+            raw_id = getattr(part.function_call, "id", None)
+            tool_call_id = str(raw_id) if raw_id else str(uuid.uuid4())
             if streaming:
                 tool_call_chunks.append(
                     tool_call_chunk(

libs/genai/langchain_google_genai/embeddings.py

@@ -448,7 +448,7 @@ def embed_documents(
             try:
                 result = self.client.models.embed_content(
                     model=self.model,
-                    contents=batch,
+                    contents=[{"parts": [{"text": text}]} for text in batch],
                     config=config,
                 )
             except ClientError as e:
@@ -566,7 +566,7 @@ async def aembed_documents(
             try:
                 result = await self.client.aio.models.embed_content(
                     model=self.model,
-                    contents=batch,
+                    contents=[{"parts": [{"text": text}]} for text in batch],
                     config=config,
                 )
             except ClientError as e:

libs/genai/pyproject.toml

@@ -13,7 +13,7 @@ version = "4.2.3"
 requires-python = ">=3.10.0,<4.0.0"
 dependencies = [
     "langchain-core>=1.3.2,<2.0.0",
-    "google-genai>=1.65.0,<2.0.0",
+    "google-genai>=1.65.0,<3.0.0",
     "pydantic>=2.0.0,<3.0.0",
     "filetype>=1.2.0,<2.0.0",
 ]

libs/genai/tests/integration_tests/test_chat_models.py

@@ -249,10 +249,15 @@ def test_chat_google_genai_invoke_with_image(backend_config: dict) -> None:
             break
     assert isinstance(result, AIMessage)
     assert isinstance(result.content, list)
-    assert isinstance(result.content[0], str)
+    if isinstance(result.content[0], dict):
+        assert result.content[0].get("type") == "text"
+        assert not result.content[0].get("text", "").startswith(" ")
+    else:
+        assert isinstance(result.content[0], str)
+        assert not result.content[0].startswith(" ")
+
     assert isinstance(result.content[1], dict)
     assert result.content[1].get("type") == "image_url"
-    assert not result.content[0].startswith(" ")
     _check_usage_metadata(result)
 
     # Test we can pass back in
@@ -276,7 +281,6 @@ def test_chat_google_genai_invoke_with_audio(backend_config: dict) -> None:
     """Test generating audio."""
     # Skip on Vertex AI - having some issues possibly upstream
     # TODO: look later
-    # https://discuss.ai.google.dev/t/request-allowlist-access-for-audio-output-in-gemini-2-5-pro-flash-tts-vertex-ai/108067
     if backend_config.get("vertexai"):
         pytest.skip("Gemini TTS on Vertex AI requires allowlist access")
 
@@ -644,15 +648,24 @@ def test_chat_google_genai_invoke_thinking_disabled(backend_config: dict) -> Non
     """Test invoking a thinking model with zero `thinking_budget`."""
     # Note certain models may not allow `thinking_budget=0`
     llm = ChatGoogleGenerativeAI(
-        model="gemini-2.5-flash", thinking_budget=0, **backend_config
+        model="gemini-3-flash-preview", thinking_budget=0, **backend_config
     )
 
     result = llm.invoke(
         "How many O's are in Google? Please tell me how you double checked the result",
     )
 
     assert isinstance(result, AIMessage)
-    assert isinstance(result.content, str)
+
+    if isinstance(result.content, list):
+        text_content = "".join(
+            block.get("text", "")
+            for block in result.content
+            if isinstance(block, dict) and block.get("type") == "text"
+        )
+        assert len(text_content) > 0
+    else:
+        assert isinstance(result.content, str)
 
     _check_usage_metadata(result)
 
@@ -1470,7 +1483,7 @@ class SimpleModel(BaseModel):
     # Initialize with thinking disabled
     # Only certain models support disabling thinking
     llm = ChatGoogleGenerativeAI(
-        model="gemini-2.5-flash",
+        model="gemini-3-flash-preview",
         thinking_budget=0,
         include_thoughts=False,
         **backend_config,
@@ -1786,7 +1799,9 @@ class MatchResult(BaseModel):
 
 def test_search_with_googletool(backend_config: dict) -> None:
     """Test using `GoogleTool` with Google Search."""
-    llm = ChatGoogleGenerativeAI(model="models/gemini-2.5-flash", **backend_config)
+    llm = ChatGoogleGenerativeAI(
+        model="models/gemini-3-flash-preview", **backend_config
+    )
     resp = llm.invoke(
         "When is the next total solar eclipse in US?",
         tools=[GoogleTool(google_search={})],
@@ -1812,7 +1827,7 @@ def test_url_context_tool(backend_config: dict) -> None:
 
 def test_google_maps_grounding(backend_config: dict) -> None:
     """Test using Google Maps grounding for location-aware responses."""
-    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", **backend_config)
+    model = ChatGoogleGenerativeAI(model="gemini-3-flash-preview", **backend_config)
     model_with_maps = model.bind_tools([{"google_maps": {}}])
 
     response = model_with_maps.invoke(
@@ -1876,7 +1891,7 @@ def test_google_maps_grounding(backend_config: dict) -> None:
 
 def test_google_maps_grounding_invoke_direct(backend_config: dict) -> None:
     """Test passing Maps grounding tool directly to invoke without binding."""
-    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", **backend_config)
+    model = ChatGoogleGenerativeAI(model="gemini-3-flash-preview", **backend_config)
 
     # Pass tools directly to invoke instead of binding
     response = model.invoke(
@@ -1994,8 +2009,7 @@ def test_chat_google_genai_invoke_with_generation_params(backend_config: dict) -
     Verifies that `max_output_tokens` (max_tokens) and `thinking_budget`
     parameters passed directly to invoke() method override model defaults.
     """
-    # Use gemini-2.5-flash because it supports thinking_budget=0
-    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", **backend_config)
+    llm = ChatGoogleGenerativeAI(model="gemini-3-flash-preview", **backend_config)
 
     # Test with max_output_tokens constraint
     result_constrained = llm.invoke(
@@ -2519,6 +2533,7 @@ def test_context_caching(backend_config: dict) -> None:
     response = chat.invoke("What is the secret number?")
 
     assert isinstance(response, AIMessage)
+
     text_blocks = [b for b in response.content_blocks if b["type"] == "text"]
     assert any("747" in b["text"] for b in text_blocks)
 

libs/genai/tests/unit_tests/test_chat_models.py

@@ -1436,6 +1436,7 @@ def test_max_retries_parameter_handling(
                                     },
                                     "grounding_chunk_indices": [0],
                                     "confidence_scores": [0.95],
+                                    "rendered_parts": None,
                                 }
                             ],
                             "web_search_queries": ["test query"],
@@ -1476,6 +1477,7 @@ def test_max_retries_parameter_handling(
                         },
                         "grounding_chunk_indices": [0],
                         "confidence_scores": [0.95],
+                        "rendered_parts": None,
                     }
                 ],
                 "image_search_queries": [],
@@ -1531,6 +1533,7 @@ def test_max_retries_parameter_handling(
                                     },
                                     "grounding_chunk_indices": [0],
                                     "confidence_scores": [0.95],
+                                    "rendered_parts": None,
                                 }
                             ],
                             "web_search_queries": ["test query"],
@@ -1572,6 +1575,7 @@ def test_max_retries_parameter_handling(
                         },
                         "grounding_chunk_indices": [0],
                         "confidence_scores": [0.95],
+                        "rendered_parts": None,
                     }
                 ],
                 "image_search_queries": ["cat images"],
@@ -1660,6 +1664,7 @@ def test_grounding_metadata_to_citations_conversion() -> None:
                             },
                             "grounding_chunk_indices": [0],
                             "confidence_scores": [0.95],
+                            "rendered_parts": None,
                         },
                         {
                             "segment": {

libs/genai/tests/unit_tests/test_embeddings.py

@@ -118,7 +118,10 @@ def test_embed_documents() -> None:
         mock_embed.assert_called_once()
         call_kwargs = mock_embed.call_args.kwargs
         assert call_kwargs["model"] == MODEL_NAME
-        assert call_kwargs["contents"] == ["test text", "test text2"]
+        assert call_kwargs["contents"] == [
+            {"parts": [{"text": "test text"}]},
+            {"parts": [{"text": "test text2"}]},
+        ]
         assert call_kwargs["config"].task_type == "RETRIEVAL_DOCUMENT"
 
         # Verify the result
@@ -320,7 +323,10 @@ async def test_aembed_documents() -> None:
         mock_embed.assert_called_once()
         call_kwargs = mock_embed.call_args.kwargs
         assert call_kwargs["model"] == MODEL_NAME
-        assert call_kwargs["contents"] == ["test text", "test text2"]
+        assert call_kwargs["contents"] == [
+            {"parts": [{"text": "test text"}]},
+            {"parts": [{"text": "test text2"}]},
+        ]
         assert call_kwargs["config"].task_type == "RETRIEVAL_DOCUMENT"
 
         # Verify the result