diff --git a/integrations/google_genai/src/haystack_integrations/components/generators/google_genai/chat/chat_generator.py b/integrations/google_genai/src/haystack_integrations/components/generators/google_genai/chat/chat_generator.py index 6e689ccf3a..33ebbec54a 100644 --- a/integrations/google_genai/src/haystack_integrations/components/generators/google_genai/chat/chat_generator.py +++ b/integrations/google_genai/src/haystack_integrations/components/generators/google_genai/chat/chat_generator.py @@ -201,6 +201,7 @@ def __init__( safety_settings: list[dict[str, Any]] | None = None, streaming_callback: StreamingCallbackT | None = None, tools: ToolsType | None = None, + google_server_tools: list[dict[str, Any]] | None = None, timeout: float | None = None, max_retries: int | None = None, ) -> None: @@ -236,6 +237,29 @@ def __init__( :param streaming_callback: A callback function that is called when a new token is received from the stream. :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls. Each tool should have a unique name. + :param google_server_tools: A list of Google server-side (built-in) tools passed directly to the API. + Each entry must be a dict matching the ``google.genai.types.Tool`` schema. The dict keys correspond + to the built-in tool type and the values are their configuration dicts (often empty ``{}``). + These are merged with any Haystack ``tools`` at request time. + + Supported built-in tools: + + - ``{"google_search": {}}`` — enables real-time Google Search grounding + - ``{"code_execution": {}}`` — enables Python code execution in a sandbox + - ``{"url_context": {}}`` — enables fetching and using URL content in context + + Example:: + + from haystack.dataclasses import ChatMessage + from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator + + chat_generator = GoogleGenAIChatGenerator( + google_server_tools=[{"google_search": {}}, {"code_execution": {}}] + ) + response = chat_generator.run([ChatMessage.from_user("What happened in AI today?")]) + + For the full list of built-in tools and their parameters, see: + https://ai.google.dev/gemini-api/docs/tools#built-in-tools :param timeout: Timeout for Google GenAI client calls. If not set, it defaults to the default set by the Google GenAI client. @@ -263,6 +287,7 @@ def __init__( self._safety_settings = safety_settings or [] self._streaming_callback = streaming_callback self._tools = tools + self._google_server_tools = google_server_tools self._timeout = timeout self._max_retries = max_retries @@ -291,6 +316,7 @@ def to_dict(self) -> dict[str, Any]: safety_settings=self._safety_settings, streaming_callback=callback_name, tools=serialized_tools, + google_server_tools=self._google_server_tools, timeout=self._timeout, max_retries=self._max_retries, ) @@ -448,9 +474,12 @@ def run( if safety_settings: config_params["safety_settings"] = safety_settings - # Add tools if provided - if tools: - config_params["tools"] = _convert_tools_to_google_genai_format(tools) + # Merge Haystack tools and Google server-side built-in tools + all_tools = _convert_tools_to_google_genai_format(tools) if tools else [] + if self._google_server_tools: + all_tools = all_tools + [types.Tool.model_validate(t) for t in self._google_server_tools] + if all_tools: + config_params["tools"] = all_tools config = types.GenerateContentConfig(**config_params) if config_params else None @@ -559,9 +588,12 @@ async def run_async( if safety_settings: config_params["safety_settings"] = safety_settings - # Add tools if provided - if tools: - config_params["tools"] = _convert_tools_to_google_genai_format(tools) + # Merge Haystack tools and Google server-side built-in tools + all_tools = _convert_tools_to_google_genai_format(tools) if tools else [] + if self._google_server_tools: + all_tools = all_tools + [types.Tool.model_validate(t) for t in self._google_server_tools] + if all_tools: + config_params["tools"] = all_tools config = types.GenerateContentConfig(**config_params) if config_params else None diff --git a/integrations/google_genai/tests/test_chat_generator.py b/integrations/google_genai/tests/test_chat_generator.py index e2df977cae..f1546e31ac 100644 --- a/integrations/google_genai/tests/test_chat_generator.py +++ b/integrations/google_genai/tests/test_chat_generator.py @@ -253,6 +253,33 @@ def test_serde_with_mixed_tools_and_toolsets(self, monkeypatch): assert restored._tools[0].name == "tool1" assert len(restored._tools[1]) == 1 + def test_init_with_google_server_tools(self, monkeypatch): + monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key") + server_tool = {"google_search": {}} + component = GoogleGenAIChatGenerator(google_server_tools=[server_tool]) + assert component._google_server_tools == [server_tool] + + def test_to_dict_with_google_server_tools(self, monkeypatch): + monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key") + server_tool = {"google_search": {}} + component = GoogleGenAIChatGenerator(google_server_tools=[server_tool]) + data = component.to_dict() + assert data["init_parameters"]["google_server_tools"] == [server_tool] + + def test_from_dict_with_google_server_tools(self, monkeypatch): + monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key") + server_tool = {"google_search": {}} + component = GoogleGenAIChatGenerator(google_server_tools=[server_tool]) + data = component.to_dict() + restored = GoogleGenAIChatGenerator.from_dict(data) + assert restored._google_server_tools == [server_tool] + + def test_to_dict_default_google_server_tools_is_none(self, monkeypatch): + monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key") + component = GoogleGenAIChatGenerator() + data = component.to_dict() + assert data["init_parameters"]["google_server_tools"] is None + def test_to_dict_with_response_format_pydantic(self, monkeypatch): """Test that to_dict serializes a Pydantic response_format to a JSON schema dict.""" monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key") @@ -365,6 +392,34 @@ def test_run_with_tools_passes_tools_to_config(self, monkeypatch, mock_response, config = call_kwargs.kwargs.get("config") or call_kwargs[1].get("config") assert config.tools is not None + def test_run_with_google_server_tools(self, monkeypatch, mock_response): + monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key") + server_tool = {"google_search": {}} + component = GoogleGenAIChatGenerator(google_server_tools=[server_tool]) + component._client.models.generate_content = Mock(return_value=mock_response) + + component.run([ChatMessage.from_user("Search for AI news")]) + + call_kwargs = component._client.models.generate_content.call_args + config = call_kwargs.kwargs.get("config") or call_kwargs[1].get("config") + assert config.tools is not None + assert len(config.tools) == 1 + + def test_run_merges_haystack_and_google_server_tools(self, monkeypatch, mock_response, tools): + monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key") + server_tool = {"google_search": {}} + component = GoogleGenAIChatGenerator(google_server_tools=[server_tool]) + component._client.models.generate_content = Mock(return_value=mock_response) + + component.run([ChatMessage.from_user("Search for AI news")], tools=tools) + + call_kwargs = component._client.models.generate_content.call_args + config = call_kwargs.kwargs.get("config") or call_kwargs[1].get("config") + assert config.tools is not None + # _convert_tools_to_google_genai_format bundles all function declarations into one Tool object, + # plus one Tool object for google_search + assert len(config.tools) == 2 + def test_run_with_safety_settings(self, monkeypatch, mock_response): monkeypatch.setenv("GOOGLE_API_KEY", "test-api-key") safety = [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}] @@ -726,6 +781,19 @@ def test_live_run_with_parallel_tools(self, tools): # Check that the response mentions both temperature readings assert "22" in final_message.text and "15" in final_message.text + def test_live_run_with_google_server_tools(self): + """ + Integration test that google_server_tools (Google Search) is accepted by the API and produces a response. + """ + component = GoogleGenAIChatGenerator(google_server_tools=[{"google_search": {}}]) + results = component.run([ChatMessage.from_user("What is the capital of France?")]) + + assert len(results["replies"]) == 1 + message = results["replies"][0] + assert message.text + assert message.meta["model"] + assert message.meta["finish_reason"] == "stop" + def test_live_run_with_thinking(self): """ Integration test for the thinking feature with a model that supports it.