From ea04ecac42cb028b6eb8ab151f80627c102b9270 Mon Sep 17 00:00:00 2001 From: james-pplx Date: Thu, 30 Apr 2026 14:51:41 +0000 Subject: [PATCH] feat(tools): add Perplexity Search tool Add a new built-in tool, PerplexitySearchTool, that calls the Perplexity Search API (POST https://api.perplexity.ai/search) via httpx.AsyncClient. The tool wraps a single `query` argument for the model and lets developers pin server-side options (max_results, recency, domain filter, etc.) at construction time, mirroring the pattern of DiscoveryEngineSearchTool. Every outgoing request includes the X-Pplx-Integration header set to google-adk/ for attribution. The API key can be supplied via constructor or PERPLEXITY_API_KEY. - src/google/adk/tools/perplexity_search_tool.py: tool implementation - src/google/adk/tools/__init__.py: lazy export - tests/unittests/tools/test_perplexity_search_tool.py: 12 tests with a mocked httpx transport, including an assertion that the attribution header is sent - contributing/samples/perplexity_search_agent: example agent Docs: https://docs.perplexity.ai/api-reference/search-post Signed-off-by: james-pplx --- .../perplexity_search_agent/__init__.py | 15 ++ .../samples/perplexity_search_agent/agent.py | 39 +++ src/google/adk/tools/__init__.py | 5 + .../adk/tools/perplexity_search_tool.py | 203 +++++++++++++++ .../tools/test_perplexity_search_tool.py | 241 ++++++++++++++++++ 5 files changed, 503 insertions(+) create mode 100644 contributing/samples/perplexity_search_agent/__init__.py create mode 100644 contributing/samples/perplexity_search_agent/agent.py create mode 100644 src/google/adk/tools/perplexity_search_tool.py create mode 100644 tests/unittests/tools/test_perplexity_search_tool.py diff --git a/contributing/samples/perplexity_search_agent/__init__.py b/contributing/samples/perplexity_search_agent/__init__.py new file mode 100644 index 0000000000..4015e47d6e --- /dev/null +++ b/contributing/samples/perplexity_search_agent/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from . import agent diff --git a/contributing/samples/perplexity_search_agent/agent.py b/contributing/samples/perplexity_search_agent/agent.py new file mode 100644 index 0000000000..943804cabf --- /dev/null +++ b/contributing/samples/perplexity_search_agent/agent.py @@ -0,0 +1,39 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sample agent using the Perplexity Search tool. + +Set the PERPLEXITY_API_KEY environment variable before running this agent. +See https://docs.perplexity.ai/api-reference/search-post for API details. +""" + +from google.adk import Agent +from google.adk.tools.perplexity_search_tool import PerplexitySearchTool + +perplexity_search = PerplexitySearchTool(max_results=5) + +root_agent = Agent( + model='gemini-2.5-flash', + name='root_agent', + description=( + 'an agent whose job it is to answer questions by searching the web' + ' via the Perplexity Search API.' + ), + instruction=( + 'You are an agent whose job is to answer questions by searching the' + ' web with the perplexity_search tool. Cite the URLs of the sources' + ' you used in your final answer.' + ), + tools=[perplexity_search], +) diff --git a/src/google/adk/tools/__init__.py b/src/google/adk/tools/__init__.py index b444e3a744..737926301f 100644 --- a/src/google/adk/tools/__init__.py +++ b/src/google/adk/tools/__init__.py @@ -36,6 +36,7 @@ from .load_artifacts_tool import load_artifacts_tool as load_artifacts from .load_memory_tool import load_memory_tool as load_memory from .long_running_tool import LongRunningFunctionTool + from .perplexity_search_tool import PerplexitySearchTool from .preload_memory_tool import preload_memory_tool as preload_memory from .tool_context import ToolContext from .transfer_to_agent_tool import transfer_to_agent @@ -79,6 +80,10 @@ '.long_running_tool', 'LongRunningFunctionTool', ), + 'PerplexitySearchTool': ( + '.perplexity_search_tool', + 'PerplexitySearchTool', + ), 'preload_memory': ('.preload_memory_tool', 'preload_memory_tool'), 'ToolContext': ('.tool_context', 'ToolContext'), 'transfer_to_agent': ('.transfer_to_agent_tool', 'transfer_to_agent'), diff --git a/src/google/adk/tools/perplexity_search_tool.py b/src/google/adk/tools/perplexity_search_tool.py new file mode 100644 index 0000000000..fdefc53e66 --- /dev/null +++ b/src/google/adk/tools/perplexity_search_tool.py @@ -0,0 +1,203 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import importlib.metadata +import logging +import os +from typing import Any +from typing import Optional + +import httpx + +from ..version import __version__ as _ADK_VERSION +from .function_tool import FunctionTool + +logger = logging.getLogger('google_adk.' + __name__) + +_PERPLEXITY_SEARCH_URL = 'https://api.perplexity.ai/search' +_INTEGRATION_SLUG = 'google-adk' +_PACKAGE_NAME = 'google-adk' +_DEFAULT_TIMEOUT_SECONDS = 30.0 + + +def _resolve_package_version() -> str: + """Returns the installed ADK package version, falling back to the in-tree version.""" + try: + return importlib.metadata.version(_PACKAGE_NAME) + except importlib.metadata.PackageNotFoundError: + return _ADK_VERSION + + +class PerplexitySearchTool(FunctionTool): + """Tool that performs web search via the Perplexity Search API. + + This tool wraps the `POST https://api.perplexity.ai/search` endpoint and + exposes a single `query` argument to the model, while letting the developer + pin server-side options (recency, domain filter, max results, etc.) at + construction time. + + See https://docs.perplexity.ai/api-reference/search-post for the request + and response schema. + + Example: + + ```python + from google.adk.agents import LlmAgent + from google.adk.tools.perplexity_search_tool import PerplexitySearchTool + + perplexity_search = PerplexitySearchTool() + agent = LlmAgent( + model='gemini-2.5-flash', + name='research_agent', + tools=[perplexity_search], + ) + ``` + """ + + def __init__( + self, + api_key: Optional[str] = None, + *, + max_results: Optional[int] = None, + max_tokens_per_page: Optional[int] = None, + country: Optional[str] = None, + search_recency_filter: Optional[str] = None, + search_domain_filter: Optional[list[str]] = None, + search_language_filter: Optional[list[str]] = None, + last_updated_after_filter: Optional[str] = None, + last_updated_before_filter: Optional[str] = None, + search_after_date_filter: Optional[str] = None, + search_before_date_filter: Optional[str] = None, + timeout: float = _DEFAULT_TIMEOUT_SECONDS, + ): + """Initializes the PerplexitySearchTool. + + Args: + api_key: The Perplexity API key. If not provided, the value of the + `PERPLEXITY_API_KEY` environment variable is used. + max_results: Maximum number of results to return (1-20). Defaults to the + API default (10) when not set. + max_tokens_per_page: Maximum tokens per page (1-1,000,000). Defaults to + the API default (4096) when not set. + country: Optional ISO 3166-1 alpha-2 country code for localization. + search_recency_filter: Optional recency filter. One of `hour`, `day`, + `week`, `month`, or `year`. + search_domain_filter: Optional list of up to 20 domains to restrict the + search to. + search_language_filter: Optional list of ISO 639-1 language codes. + last_updated_after_filter: Optional `MM/DD/YYYY` lower bound on the + `last_updated` field of results. + last_updated_before_filter: Optional `MM/DD/YYYY` upper bound on the + `last_updated` field of results. + search_after_date_filter: Optional `MM/DD/YYYY` lower bound on the + result publication date. + search_before_date_filter: Optional `MM/DD/YYYY` upper bound on the + result publication date. + timeout: HTTP timeout in seconds for each search request. + + Raises: + ValueError: If no API key is supplied and `PERPLEXITY_API_KEY` is not + set in the environment. + """ + super().__init__(self.perplexity_search) + resolved_api_key = api_key or os.environ.get('PERPLEXITY_API_KEY') + if not resolved_api_key: + raise ValueError( + 'Perplexity API key is required: pass `api_key` to ' + 'PerplexitySearchTool or set the PERPLEXITY_API_KEY ' + 'environment variable.' + ) + self._api_key = resolved_api_key + self._max_results = max_results + self._max_tokens_per_page = max_tokens_per_page + self._country = country + self._search_recency_filter = search_recency_filter + self._search_domain_filter = search_domain_filter + self._search_language_filter = search_language_filter + self._last_updated_after_filter = last_updated_after_filter + self._last_updated_before_filter = last_updated_before_filter + self._search_after_date_filter = search_after_date_filter + self._search_before_date_filter = search_before_date_filter + self._timeout = timeout + + def _build_headers(self) -> dict[str, str]: + return { + 'Authorization': f'Bearer {self._api_key}', + 'Content-Type': 'application/json', + 'X-Pplx-Integration': ( + f'{_INTEGRATION_SLUG}/{_resolve_package_version()}' + ), + } + + def _build_body(self, query: str) -> dict[str, Any]: + body: dict[str, Any] = {'query': query} + optional_fields: dict[str, Any] = { + 'max_results': self._max_results, + 'max_tokens_per_page': self._max_tokens_per_page, + 'country': self._country, + 'search_recency_filter': self._search_recency_filter, + 'search_domain_filter': self._search_domain_filter, + 'search_language_filter': self._search_language_filter, + 'last_updated_after_filter': self._last_updated_after_filter, + 'last_updated_before_filter': self._last_updated_before_filter, + 'search_after_date_filter': self._search_after_date_filter, + 'search_before_date_filter': self._search_before_date_filter, + } + for key, value in optional_fields.items(): + if value is not None: + body[key] = value + return body + + async def perplexity_search(self, query: str) -> dict[str, Any]: + """Searches the web via the Perplexity Search API. + + Args: + query: The search query. + + Returns: + A dictionary with `status` set to either `success` or `error`. On + success, `results` contains a list of result entries with `title`, + `url`, `snippet`, `date`, and `last_updated` fields, along with the + raw `id` and `server_time` returned by the API. + """ + headers = self._build_headers() + body = self._build_body(query) + + try: + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.post( + _PERPLEXITY_SEARCH_URL, headers=headers, json=body + ) + response.raise_for_status() + payload = response.json() + except httpx.HTTPStatusError as e: + logger.exception('Perplexity Search request failed.') + return { + 'status': 'error', + 'error_message': ( + f'Perplexity Search returned HTTP {e.response.status_code}.' + ), + } + except httpx.HTTPError as e: + logger.exception('Perplexity Search request failed.') + return {'status': 'error', 'error_message': str(e)} + + return { + 'status': 'success', + 'results': payload.get('results', []), + 'id': payload.get('id'), + 'server_time': payload.get('server_time'), + } diff --git a/tests/unittests/tools/test_perplexity_search_tool.py b/tests/unittests/tools/test_perplexity_search_tool.py new file mode 100644 index 0000000000..1be5e2e35a --- /dev/null +++ b/tests/unittests/tools/test_perplexity_search_tool.py @@ -0,0 +1,241 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +from typing import Any + +from google.adk.tools.perplexity_search_tool import PerplexitySearchTool +import httpx +import pytest + +_FAKE_PAYLOAD: dict[str, Any] = { + "results": [{ + "title": "Example", + "url": "https://example.com", + "snippet": "An example result.", + "date": "2026-01-01", + "last_updated": "2026-01-02", + }], + "id": "search-id-123", + "server_time": "2026-04-30T00:00:00Z", +} + + +class _CapturedRequest: + """Holds the most recent request seen by the fake transport.""" + + def __init__(self) -> None: + self.url: str | None = None + self.headers: httpx.Headers | None = None + self.json: dict[str, Any] | None = None + self.method: str | None = None + + +def _make_transport( + *, + captured: _CapturedRequest, + payload: dict[str, Any] | None = None, + status_code: int = 200, + raise_exc: Exception | None = None, +) -> httpx.MockTransport: + """Builds a MockTransport that captures the outgoing request and returns a fixed response.""" + + def handler(request: httpx.Request) -> httpx.Response: + captured.method = request.method + captured.url = str(request.url) + captured.headers = request.headers + captured.json = json.loads(request.content) if request.content else None + if raise_exc is not None: + raise raise_exc + return httpx.Response(status_code, json=payload or {}) + + return httpx.MockTransport(handler) + + +@pytest.fixture(autouse=True) +def _patch_async_client(monkeypatch): + """Routes every httpx.AsyncClient created by the tool through a MockTransport.""" + + state: dict[str, Any] = { + "captured": _CapturedRequest(), + "payload": _FAKE_PAYLOAD, + "status_code": 200, + "raise_exc": None, + } + + original_init = httpx.AsyncClient.__init__ + + def _patched_init(self, *args, **kwargs): + kwargs["transport"] = _make_transport( + captured=state["captured"], + payload=state["payload"], + status_code=state["status_code"], + raise_exc=state["raise_exc"], + ) + return original_init(self, *args, **kwargs) + + monkeypatch.setattr(httpx.AsyncClient, "__init__", _patched_init) + return state + + +@pytest.mark.asyncio +async def test_init_requires_api_key(monkeypatch): + monkeypatch.delenv("PERPLEXITY_API_KEY", raising=False) + with pytest.raises(ValueError, match="Perplexity API key"): + PerplexitySearchTool() + + +@pytest.mark.asyncio +async def test_init_reads_api_key_from_environment(monkeypatch): + monkeypatch.setenv("PERPLEXITY_API_KEY", "env-key") + tool = PerplexitySearchTool() + assert tool._api_key == "env-key" + + +@pytest.mark.asyncio +async def test_init_prefers_explicit_api_key(monkeypatch): + monkeypatch.setenv("PERPLEXITY_API_KEY", "env-key") + tool = PerplexitySearchTool(api_key="explicit-key") + assert tool._api_key == "explicit-key" + + +@pytest.mark.asyncio +async def test_search_sends_attribution_header(_patch_async_client): + tool = PerplexitySearchTool(api_key="test-key") + result = await tool.perplexity_search("hello world") + + captured = _patch_async_client["captured"] + assert result["status"] == "success" + assert captured.headers is not None + integration_header = captured.headers.get("x-pplx-integration") + assert integration_header is not None + assert integration_header.startswith("google-adk/") + + +@pytest.mark.asyncio +async def test_search_sends_bearer_auth_and_json_content_type( + _patch_async_client, +): + tool = PerplexitySearchTool(api_key="test-key") + await tool.perplexity_search("foo") + + captured = _patch_async_client["captured"] + assert captured.headers["authorization"] == "Bearer test-key" + assert captured.headers["content-type"] == "application/json" + assert captured.method == "POST" + assert captured.url == "https://api.perplexity.ai/search" + + +@pytest.mark.asyncio +async def test_search_returns_results_payload(_patch_async_client): + tool = PerplexitySearchTool(api_key="test-key") + result = await tool.perplexity_search("hello") + + assert result["status"] == "success" + assert result["results"] == _FAKE_PAYLOAD["results"] + assert result["id"] == "search-id-123" + assert result["server_time"] == "2026-04-30T00:00:00Z" + + +@pytest.mark.asyncio +async def test_search_body_contains_only_query_when_no_options( + _patch_async_client, +): + tool = PerplexitySearchTool(api_key="test-key") + await tool.perplexity_search("hello world") + + captured = _patch_async_client["captured"] + assert captured.json == {"query": "hello world"} + + +@pytest.mark.asyncio +async def test_search_body_includes_configured_options(_patch_async_client): + tool = PerplexitySearchTool( + api_key="test-key", + max_results=5, + max_tokens_per_page=1024, + country="US", + search_recency_filter="week", + search_domain_filter=["example.com", "wikipedia.org"], + search_language_filter=["en"], + last_updated_after_filter="01/01/2026", + last_updated_before_filter="04/30/2026", + search_after_date_filter="01/01/2025", + search_before_date_filter="12/31/2025", + ) + await tool.perplexity_search("topic") + + body = _patch_async_client["captured"].json + assert body == { + "query": "topic", + "max_results": 5, + "max_tokens_per_page": 1024, + "country": "US", + "search_recency_filter": "week", + "search_domain_filter": ["example.com", "wikipedia.org"], + "search_language_filter": ["en"], + "last_updated_after_filter": "01/01/2026", + "last_updated_before_filter": "04/30/2026", + "search_after_date_filter": "01/01/2025", + "search_before_date_filter": "12/31/2025", + } + + +@pytest.mark.asyncio +async def test_search_returns_error_on_http_status(_patch_async_client): + _patch_async_client["status_code"] = 401 + _patch_async_client["payload"] = {"error": "unauthorized"} + + tool = PerplexitySearchTool(api_key="bad-key") + result = await tool.perplexity_search("anything") + + assert result["status"] == "error" + assert "401" in result["error_message"] + + +@pytest.mark.asyncio +async def test_search_returns_error_on_transport_error(_patch_async_client): + _patch_async_client["raise_exc"] = httpx.ConnectError("boom") + + tool = PerplexitySearchTool(api_key="test-key") + result = await tool.perplexity_search("anything") + + assert result["status"] == "error" + assert "boom" in result["error_message"] + + +@pytest.mark.asyncio +async def test_tool_function_declaration_uses_query_argument( + _patch_async_client, +): + tool = PerplexitySearchTool(api_key="test-key") + declaration = tool._get_declaration() + + assert declaration is not None + assert declaration.name == "perplexity_search" + assert declaration.parameters is not None + properties = declaration.parameters.properties or {} + assert "query" in properties + + +@pytest.mark.asyncio +async def test_tool_is_lazily_exported(monkeypatch): + monkeypatch.setenv("PERPLEXITY_API_KEY", "env-key") + from google.adk import tools as adk_tools + + tool_cls = adk_tools.PerplexitySearchTool + instance = tool_cls() + assert instance.name == "perplexity_search"