diff --git a/docs/mkdocs/zh/tool.md b/docs/mkdocs/zh/tool.md index b7e2cd1d..ff265cc5 100644 --- a/docs/mkdocs/zh/tool.md +++ b/docs/mkdocs/zh/tool.md @@ -2268,21 +2268,32 @@ print(regular_tool.is_streaming) # False ## WebFetchTool (网页获取工具) -`WebFetchTool` 是 trpc-agent-python 框架内置的**单 URL 联网抓取工具**。当 Agent 需要阅读、摘要或引用某个公开网页的内容时，可以通过该工具发起一次 HTTP GET 请求，框架会将响应统一转换为可供 LLM 消费的结构化文本：HTML 会被裁剪为 Markdown 纯文本，其它 `text/*` / `application/json` 等文本型 MIME 按原样返回，二进制响应则以结构化错误拒收。 +`WebFetchTool` 是 trpc-agent-python 框架内置的**单 URL 联网抓取工具**。当 Agent 需要阅读、摘要或引用某个公开网页的内容时，可以通过该工具获取页面文本，并统一转换为可供 LLM 消费的结构化结果。 + +该工具采用**可插拔 provider** 设计，目前内置两种后端： + +- **`direct`（默认）**：直接发起 HTTP GET。HTML 会被裁剪为 Markdown 纯文本，其它 `text/*` / `application/json` 等文本型 MIME 按原样返回，二进制响应以结构化错误拒收 +- **`tavily`**：走 Tavily Extract API，由 Tavily 负责抽取页面正文；需要配置 `api_key`（或环境变量 `TAVILY_API_KEY`），适合希望减少本地 HTML 解析噪音、直接拿到较干净正文的场景 ### 功能特性 -- **单次 HTTP GET**：HTML 自动转换为 Markdown 纯文本（去除 `" + client = _make_mock_client(get_responses={ + "/": (200, html, { + "content-type": "text/html; charset=utf-8" + }), + }) + tool = WebFetchTool( + http_client=client, + block_private_network=False, + max_content_length=10_000, + ) + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "https://example.com/"}, + ) + assert res["error"] == "" + assert res["status_code"] == 200 + assert res["content_type"] == "text/html" + assert "# Hello" in res["content"] + assert "World" in res["content"] + assert "evil()" not in res["content"] + await client.aclose() + + @pytest.mark.asyncio + async def test_max_length_truncates_content(self): + client = _make_mock_client(get_responses={ + "/": (200, b"ABCDEFGHIJKLMNOPQRSTUVWXYZ", { + "content-type": "text/plain" + }), + }) + tool = WebFetchTool( + http_client=client, + block_private_network=False, + max_content_length=100, + ) + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={ + "url": "https://example.com/", + "max_length": 10, + }, + ) + assert res["content"].endswith("...") + assert len(res["content"]) <= 13 + await client.aclose() + + @pytest.mark.asyncio + async def test_cache_hit_marks_cached_true(self): + client = _make_mock_client(get_responses={ + "/": (200, b"cached body", { + "content-type": "text/plain" + }), + }) + tool = WebFetchTool( + http_client=client, + block_private_network=False, + enable_cache=True, + cache_ttl_seconds=60.0, + cache_max_bytes=1024 * 1024, + ) + first = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "https://example.com/"}, + ) + second = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "https://www.example.com/"}, + ) + assert first["cached"] is False + assert second["cached"] is True + assert second["content"] == first["content"] + assert len(client._captured["all_requests"]) == 1 + await client.aclose() + + +_TAVILY_EXTRACT_RESPONSE: Dict[str, Any] = { + "results": [{ + "url": "https://docs.python.org/3/whatsnew/3.13.html", + "raw_content": "# What's New In Python 3.13\n\nFree-threaded CPython.", + }], +} + + +class TestTavilyProvider: + + @pytest.mark.asyncio + async def test_missing_credentials_returns_structured_error(self, monkeypatch): + monkeypatch.delenv("TAVILY_API_KEY", raising=False) + tool = WebFetchTool(provider="tavily", api_key="") + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "https://example.com"}, + ) + assert "TAVILY_NOT_CONFIGURED" in res["error"] + + @pytest.mark.asyncio + async def test_happy_path_posts_extract_payload(self): + client = _make_mock_client(post_json={"/extract": _TAVILY_EXTRACT_RESPONSE}) + tool = WebFetchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/extract", + block_private_network=False, + tavily_extract_depth="advanced", + tavily_extra_params={"include_images": True}, + ) + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "https://docs.python.org/3/whatsnew/3.13.html"}, + ) + assert res["error"] == "" + assert res["status_code"] == 200 + assert res["content_type"] == "text/markdown" + assert res["url"] == "https://docs.python.org/3/whatsnew/3.13.html" + assert "Free-threaded CPython" in res["content"] + + req = client._captured["last_request"] + assert req.method == "POST" + assert req.headers.get("authorization") == "Bearer tvly-test" + payload = json.loads(req.content) + assert payload["urls"] == ["https://docs.python.org/3/whatsnew/3.13.html"] + assert payload["extract_depth"] == "advanced" + assert payload["include_images"] is True + await client.aclose() + + @pytest.mark.asyncio + async def test_failed_extract_returns_structured_error(self): + client = _make_mock_client( + post_json={"/extract": { + "results": [], + "failed_results": ["timeout while extracting"], + }}) + tool = WebFetchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/extract", + block_private_network=False, + ) + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "https://example.com/missing"}, + ) + assert "TAVILY_EXTRACT_ERROR" in res["error"] + assert "timeout" in res["error"] + await client.aclose() + + @pytest.mark.asyncio + async def test_http_status_error_maps_to_http_error(self): + client = _make_mock_client( + post_json={"/extract": { + "error": "unauthorized" + }}, + post_status={"/extract": 401}, + ) + tool = WebFetchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/extract", + block_private_network=False, + ) + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "https://example.com/page"}, + ) + assert "HTTP_ERROR" in res["error"] + assert "401" in res["error"] + assert res["status_code"] == 0 + await client.aclose() + + @pytest.mark.asyncio + async def test_ssrf_guard_blocks_private_targets_before_tavily(self): + client = _make_mock_client(post_json={"/extract": _TAVILY_EXTRACT_RESPONSE}) + tool = WebFetchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/extract", + block_private_network=True, + ) + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "http://127.0.0.1:8080/secret"}, + ) + assert "SSRF_BLOCKED_URL" in res["error"] + assert client._captured["all_requests"] == [] + await client.aclose() + + @pytest.mark.asyncio + async def test_content_fallback_uses_content_field(self): + client = _make_mock_client(post_json={ + "/extract": { + "results": [{ + "url": "https://example.com/page", + "content": "plain content fallback", + }], + } + }) + tool = WebFetchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/extract", + block_private_network=False, + ) + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={"url": "https://example.com/page"}, + ) + assert res["content"] == "plain content fallback" + await client.aclose() + + @pytest.mark.asyncio + async def test_max_length_truncates_tavily_content(self): + client = _make_mock_client(post_json={ + "/extract": { + "results": [{ + "url": "https://example.com/page", + "raw_content": "ABCDEFGHIJKLMNOPQRSTUVWXYZ", + }], + } + }) + tool = WebFetchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/extract", + block_private_network=False, + ) + res = await tool._run_async_impl( + tool_context=_tool_ctx(), + args={ + "url": "https://example.com/page", + "max_length": 8, + }, + ) + assert res["content"].endswith("...") + assert len(res["content"]) <= 11 + await client.aclose() + + +class TestHelpers: + + def test_truncate(self): + assert _truncate("abc", 10) == "abc" + assert _truncate("x" * 20, 10) == ("x" * 10) + "..." + assert _truncate("keep-all", 0) == "keep-all" + + def test_normalise_content_type(self): + assert _normalise_content_type("text/html; charset=utf-8") == "text/html" + assert _normalise_content_type("") == "" + + def test_cache_key_normalises_variants(self): + assert _cache_key("https://www.Example.com/path/") == _cache_key("https://example.com/path") + + def test_is_blocked_url(self): + assert _is_blocked_url("https://docs.python.org/x", None, ["python.org"]) + assert not _is_blocked_url("https://python.org/x", ["python.org"], None) + assert _is_blocked_url("https://example.com/x", ["python.org"], None) diff --git a/tests/tools/test_websearch_tool.py b/tests/tools/test_websearch_tool.py index 518fa2a2..014310c1 100644 --- a/tests/tools/test_websearch_tool.py +++ b/tests/tools/test_websearch_tool.py @@ -6,9 +6,9 @@ """Tests for :mod:`trpc_agent_sdk.tools._websearch_tool`. Covers the full BaseTool surface area: -- ``SearchHit`` / ``WebSearchResult`` schemas +- ``SearchHit`` / ``WebSearchResult`` / ``TavilyWebSearchResult`` schemas - Constructor validation and provider dispatch -- ``_get_declaration`` parameter shape +- ``_get_declaration`` parameter shape (including Tavily ``include_images``) - Input validation (missing query, short query, mutually exclusive ``allowed_domains`` / ``blocked_domains``) - DuckDuckGo Instant Answer path (summary aggregation, related topics, @@ -16,16 +16,21 @@ - Google CSE path (items, pagemap metatag enrichment, spell-corrected effective query, server-side domain filter parameters, API error, missing credentials) +- Tavily Search path (results, optional images, domain include/exclude, + API error, missing credentials, extra params) - HTTP errors surfacing as structured tool errors - ``process_request`` registering the declaration and appending the "current month / Sources" system instruction - Internal helpers (``_truncate``, ``_extract_domain_from_url``, ``_is_blocked``, ``_extract_title_from_ddg_topic``, ``_extract_desc_from_pagemap``) -- Module-level singleton + auto-registration with ``ToolRegistry`` """ from __future__ import annotations +# Ensure ``pydantic.root_model`` is registered before MCP/google-genai import +# chains run; otherwise collection can hit KeyError: 'pydantic.root_model'. +import pydantic.root_model # noqa: F401 + import json from typing import Any from typing import Dict @@ -36,7 +41,9 @@ from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.models import LlmRequest +from trpc_agent_sdk.tools._websearch_tool import ImageHit from trpc_agent_sdk.tools._websearch_tool import SearchHit +from trpc_agent_sdk.tools._websearch_tool import TavilyWebSearchResult from trpc_agent_sdk.tools._websearch_tool import WebSearchResult from trpc_agent_sdk.tools._websearch_tool import WebSearchTool from trpc_agent_sdk.tools._websearch_tool import _current_month_year @@ -114,6 +121,24 @@ def test_model_dump_is_json_serialisable(self): assert parsed["results"][0]["title"] == "t" +class TestTavilyWebSearchResultSchema: + + def test_defaults_include_images(self): + r = TavilyWebSearchResult(query="q", provider="tavily") + assert r.images == [] + assert r.results == [] + + def test_roundtrip_with_images(self): + r = TavilyWebSearchResult( + query="q", + provider="tavily", + results=[SearchHit(title="t", url="https://x", snippet="s")], + images=[ImageHit(url="https://img.example/a.png", description="alt")], + ) + restored = TavilyWebSearchResult.model_validate(r.model_dump()) + assert restored == r + + class TestWebSearchToolInit: def test_default_is_duckduckgo(self): @@ -121,13 +146,9 @@ def test_default_is_duckduckgo(self): assert t.name == _TOOL_NAME assert t._provider == "duckduckgo" - def test_invalid_provider_accepted_by_construction(self): - # Provider validity is enforced only by the Literal type hint at - # static-analysis time; construction itself does not raise. - # Runtime dispatch falls into the Google branch, which returns a - # helpful error rather than crashing. - t = WebSearchTool(provider="bing") # type: ignore[arg-type] - assert t._provider == "bing" + def test_invalid_provider_raises(self): + with pytest.raises(ValueError, match="Unsupported web search provider"): + WebSearchTool(provider="bing") # type: ignore[arg-type] def test_google_without_creds_warns_not_raises(self, caplog): # Capture warnings from the tool's logger. @@ -135,6 +156,19 @@ def test_google_without_creds_warns_not_raises(self, caplog): t = WebSearchTool(provider="google", api_key="", engine_id="") assert t._provider == "google" + def test_tavily_without_creds_warns_not_raises(self, monkeypatch): + # Empty ``api_key`` falls back to ``TAVILY_API_KEY``; clear it so the + # missing-credentials path is exercised even when CI exports the var. + monkeypatch.delenv("TAVILY_API_KEY", raising=False) + t = WebSearchTool(provider="tavily", api_key="") + assert t._provider == "tavily" + assert t._api_key == "" + + def test_tavily_reads_api_key_from_env(self, monkeypatch): + monkeypatch.setenv("TAVILY_API_KEY", "env-tavily-key") + t = WebSearchTool(provider="tavily") + assert t._api_key == "env-tavily-key" + def test_results_num_clamped(self): from trpc_agent_sdk.tools._websearch_tool import _MAX_COUNT # Above the cap → clamped to _MAX_COUNT. @@ -162,6 +196,12 @@ def test_declaration_shape(self): assert props["allowed_domains"].type == Type.ARRAY assert props["allowed_domains"].items is not None + def test_tavily_declaration_exposes_include_images(self): + decl = WebSearchTool(provider="tavily", api_key="k")._get_declaration() + props = decl.parameters.properties + assert "include_images" in props + assert props["include_images"].type == Type.BOOLEAN + class TestInputValidation: @@ -1025,6 +1065,280 @@ async def test_api_error_surfaced_as_structured_result(self): await client.aclose() +_TAVILY_RESPONSE: Dict[str, Any] = { + "answer": + "Python 3.13 adds free-threaded support.", + "results": [ + { + "title": "What's New In Python 3.13", + "url": "https://docs.python.org/3/whatsnew/3.13.html", + "content": "This article explains the new features in Python 3.13.", + }, + { + "title": "Python on Wikipedia", + "url": "https://en.wikipedia.org/wiki/Python_(programming_language)", + "content": "Python is a high-level programming language.", + "images": [{ + "url": "https://upload.wikimedia.org/python.png", + "description": "Python logo", + }], + }, + ], + "images": [ + "https://cdn.example.com/python.jpg", + { + "url": "https://cdn.example.com/python.jpg", + "description": "duplicate image", + }, + { + "url": "https://images.python.org/hero.png", + "description": "Python hero image", + }, + ], +} + + +class TestTavilyProvider: + + @pytest.mark.asyncio + async def test_missing_credentials_returns_helpful_result(self, monkeypatch): + monkeypatch.delenv("TAVILY_API_KEY", raising=False) + t = WebSearchTool(provider="tavily", api_key="") + res = await t._run_async_impl( + tool_context=_tool_ctx(), + args={"query": "python"}, + ) + assert res["provider"] == "tavily" + assert res["results"] == [] + assert res["images"] == [] + assert "not configured" in res["summary"] + + @pytest.mark.asyncio + async def test_happy_path_posts_payload_and_parses_results(self): + client = _make_mock_client({"/search": _TAVILY_RESPONSE}) + t = WebSearchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/search", + ) + res = await t._run_async_impl( + tool_context=_tool_ctx(), + args={ + "query": "Python 3.13", + "count": 2, + }, + ) + assert res["provider"] == "tavily" + assert res["query"] == "Python 3.13" + # The implementation surfaces ``answer`` whenever Tavily returns it. + assert "free-threaded" in res["summary"] + assert [h["url"] for h in res["results"]] == [ + "https://docs.python.org/3/whatsnew/3.13.html", + "https://en.wikipedia.org/wiki/Python_(programming_language)", + ] + assert res["images"] == [] + + req = client._captured["last_request"] + assert req.method == "POST" + assert req.headers.get("authorization") == "Bearer tvly-test" + payload = json.loads(req.content) + assert payload["query"] == "Python 3.13" + assert payload["max_results"] == 2 + assert payload["include_images"] is False + assert payload["include_answer"] is False + await client.aclose() + + @pytest.mark.asyncio + async def test_include_answer_fills_summary(self): + client = _make_mock_client({"/search": _TAVILY_RESPONSE}) + t = WebSearchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/search", + tavily_extra_params={"include_answer": True}, + ) + res = await t._run_async_impl( + tool_context=_tool_ctx(), + args={"query": "Python 3.13"}, + ) + assert "free-threaded" in res["summary"] + payload = json.loads(client._captured["last_request"].content) + assert payload["include_answer"] is True + await client.aclose() + + @pytest.mark.asyncio + async def test_include_images_returns_deduped_image_hits(self): + client = _make_mock_client({"/search": _TAVILY_RESPONSE}) + t = WebSearchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/search", + results_num=5, + ) + res = await t._run_async_impl( + tool_context=_tool_ctx(), + args={ + "query": "python", + "include_images": True, + }, + ) + image_urls = [img["url"] for img in res["images"]] + assert image_urls == [ + "https://cdn.example.com/python.jpg", + "https://images.python.org/hero.png", + "https://upload.wikimedia.org/python.png", + ] + assert res["images"][2]["description"] == "Python logo" + + req = client._captured["last_request"] + payload = json.loads(req.content) + assert payload["include_images"] is True + await client.aclose() + + @pytest.mark.asyncio + async def test_allowed_domains_mapped_to_include_domains(self): + client = _make_mock_client({"/search": _TAVILY_RESPONSE}) + t = WebSearchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/search", + ) + res = await t._run_async_impl( + tool_context=_tool_ctx(), + args={ + "query": "python", + "allowed_domains": ["python.org"], + }, + ) + req = client._captured["last_request"] + payload = json.loads(req.content) + assert payload["include_domains"] == ["python.org"] + assert "exclude_domains" not in payload + # Client-side filter keeps only python.org (docs.python.org matches). + assert [h["url"] for h in res["results"]] == [ + "https://docs.python.org/3/whatsnew/3.13.html", + ] + await client.aclose() + + @pytest.mark.asyncio + async def test_blocked_domains_mapped_to_exclude_domains(self): + client = _make_mock_client({"/search": _TAVILY_RESPONSE}) + t = WebSearchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/search", + ) + res = await t._run_async_impl( + tool_context=_tool_ctx(), + args={ + "query": "python", + "blocked_domains": ["wikipedia.org"], + }, + ) + req = client._captured["last_request"] + payload = json.loads(req.content) + assert payload["exclude_domains"] == ["wikipedia.org"] + urls = [h["url"] for h in res["results"]] + assert "https://en.wikipedia.org/wiki/Python_(programming_language)" not in urls + assert "https://docs.python.org/3/whatsnew/3.13.html" in urls + await client.aclose() + + @pytest.mark.asyncio + async def test_results_are_deduplicated_by_url(self): + body = { + "results": [ + { + "title": "Python docs", + "url": "https://docs.python.org/3", + "content": "first", + }, + { + "title": "Python docs (dup)", + "url": "https://www.docs.python.org/3/", + "content": "duplicate", + }, + { + "title": "PEP 8", + "url": "https://peps.python.org/pep-0008/", + "content": "style", + }, + ], + } + client = _make_mock_client({"/search": body}) + t = WebSearchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/search", + ) + res = await t._run_async_impl(tool_context=_tool_ctx(), args={"query": "python"}) + assert [h["url"] for h in res["results"]] == [ + "https://docs.python.org/3", + "https://peps.python.org/pep-0008/", + ] + await client.aclose() + + @pytest.mark.asyncio + async def test_tavily_extra_params_are_merged(self): + client = _make_mock_client({"/search": {"results": []}}) + t = WebSearchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/search", + tavily_extra_params={ + "search_depth": "advanced", + "include_answer": True, + }, + ) + await t._run_async_impl(tool_context=_tool_ctx(), args={"query": "python"}) + payload = json.loads(client._captured["last_request"].content) + assert payload["search_depth"] == "advanced" + assert payload["include_answer"] is True + await client.aclose() + + @pytest.mark.asyncio + async def test_api_error_surfaced_as_structured_result(self): + client = _make_mock_client({"/search": {"error": "Invalid API key"}}) + t = WebSearchTool( + provider="tavily", + api_key="bad", + http_client=client, + base_url="https://api.tavily.com/search", + ) + res = await t._run_async_impl(tool_context=_tool_ctx(), args={"query": "python"}) + assert res["results"] == [] + assert res["images"] == [] + assert "Invalid API key" in res["summary"] + await client.aclose() + + @pytest.mark.asyncio + async def test_count_limits_hits_and_images(self): + client = _make_mock_client({"/search": _TAVILY_RESPONSE}) + t = WebSearchTool( + provider="tavily", + api_key="tvly-test", + http_client=client, + base_url="https://api.tavily.com/search", + results_num=1, + ) + res = await t._run_async_impl( + tool_context=_tool_ctx(), + args={ + "query": "python", + "include_images": True, + }, + ) + assert len(res["results"]) == 1 + assert len(res["images"]) == 1 + await client.aclose() + + class TestHttpErrorHandling: @pytest.mark.asyncio diff --git a/trpc_agent_sdk/tools/_webfetch_tool.py b/trpc_agent_sdk/tools/_webfetch_tool.py index 363b8a60..8c266b2e 100644 --- a/trpc_agent_sdk/tools/_webfetch_tool.py +++ b/trpc_agent_sdk/tools/_webfetch_tool.py @@ -7,7 +7,7 @@ Provides a client-side :class:`WebFetchTool` that: -- issues an unauthenticated HTTP GET against a single URL, +- fetches a single URL directly or through Tavily Extract, - converts the response to text (Markdown for HTML; verbatim for textual MIME types), - optionally caches raw fetch results in a TTL + byte-bounded LRU so @@ -18,12 +18,14 @@ import asyncio import ipaddress +import os import re import socket import time from collections import OrderedDict from typing import Any from typing import List +from typing import Literal from typing import Optional from typing import Tuple from urllib.parse import urlparse @@ -58,6 +60,8 @@ # Cache defaults (15 minutes / 50 MiB) _DEFAULT_CACHE_TTL_SECONDS = 15.0 * 60 _DEFAULT_CACHE_MAX_BYTES = 50 * 1024 * 1024 +# Tavily Extract base URL. +_TAVILY_BASE_URL = "https://api.tavily.com/extract" # MIME types we render as Markdown-ish text. _HTML_TYPES = frozenset(["text/html", "application/xhtml+xml"]) # MIME types we return verbatim. Any unknown ``text/*`` MIME is accepted @@ -77,14 +81,14 @@ ]) # WebFetchTool description shown to the LLM _BASE_DESCRIPTION = """\ -Fetch a single URL via HTTP GET and return its textual content. +Fetch a single public URL and return its textual content. - HTML pages are stripped to Markdown-ish plain text; other textual MIME types are returned verbatim. - Content is capped to a finite length (override via `max_length`) so large pages do not overflow the context window. - Binary content (PDF / image / archive / ...) is rejected with a structured error; use a dedicated tool for it. - The URL must be an absolute http(s) URL. The tool's domain allow/block lists are enforced before the request. - By default the tool refuses to dial loopback / private / link-local targets (e.g. 127.0.0.1, 169.254.169.254, intranet IPs) so it cannot be abused as an SSRF probe. -- This tool is read-only and does not modify any files or remote state. + - This tool is read-only and does not modify any files or remote state. USE WHEN the user asks to: - read, summarise, quote or extract a fact from a specific webpage, doc page, RFC, changelog, or news article @@ -94,7 +98,7 @@ - an MCP-provided web fetch tool is available — prefer it, as it may have fewer restrictions and richer auth - the URL requires authentication or session cookies (Google Docs, Confluence, Jira, private GitHub, intranet) — prefer a dedicated MCP/authenticated tool - - you need to render JavaScript, submit forms, or perform interactive browsing (this tool only issues one GET) + - you need to render JavaScript, submit forms, or perform interactive browsing - the content is binary — the tool rejects non-textual responses Usage notes: @@ -104,6 +108,8 @@ \ """ +FetchProviderType = Literal["direct", "tavily"] + class FetchResult(BaseModel): """Structured output of :class:`WebFetchTool`.""" @@ -372,8 +378,8 @@ async def clear(self) -> None: class WebFetchTool(BaseTool): """LLM tool that fetches a single URL and returns its textual content. - The tool issues an unauthenticated HTTP GET, converts the response - to text (Markdown-ish for HTML, verbatim for other textual types), + The tool fetches directly with an unauthenticated HTTP GET or delegates + extraction to Tavily, then converts the response to text, enforces tool-level allow/block lists, and caps the returned content to protect the model context window. Binary responses are rejected with a structured error rather than dumped as base64 blobs. @@ -385,6 +391,11 @@ class WebFetchTool(BaseTool): freshness. Args: + provider: Fetch backend, ``"direct"`` (default) or ``"tavily"``. + api_key: Tavily API key; falls back to ``TAVILY_API_KEY``. + base_url: Override the Tavily Extract endpoint for tests or proxies. + tavily_extract_depth: Tavily extraction depth, ``"basic"`` or ``"advanced"``. + tavily_extra_params: Additional Tavily Extract JSON parameters. timeout: HTTP timeout in seconds. user_agent: HTTP ``User-Agent`` header. proxy: Optional HTTP proxy URL forwarded to httpx. @@ -407,6 +418,11 @@ class WebFetchTool(BaseTool): def __init__( self, *, + provider: FetchProviderType = "direct", + api_key: Optional[str] = None, + base_url: Optional[str] = None, + tavily_extract_depth: Literal["basic", "advanced"] = "basic", + tavily_extra_params: Optional[dict[str, Any]] = None, timeout: float = _DEFAULT_TIMEOUT, user_agent: str = _DEFAULT_USER_AGENT, proxy: Optional[str] = None, @@ -424,12 +440,21 @@ def __init__( filters_name: Optional[List[str]] = None, filters: Optional[List[BaseFilter]] = None, ) -> None: + if provider not in ("direct", "tavily"): + raise ValueError(f"Unsupported web fetch provider: {provider!r}") + if tavily_extract_depth not in ("basic", "advanced"): + raise ValueError("tavily_extract_depth must be 'basic' or 'advanced'") super().__init__( name="webfetch", description=_BASE_DESCRIPTION, filters_name=filters_name, filters=filters, ) + self._provider: FetchProviderType = provider + self._api_key = api_key or os.environ.get("TAVILY_API_KEY", "") + self._base_url = base_url or _TAVILY_BASE_URL + self._tavily_extract_depth = tavily_extract_depth + self._tavily_extra_params = tavily_extra_params or {} self._timeout = float(timeout) self._user_agent = user_agent self._proxy = proxy @@ -450,6 +475,9 @@ def __init__( ttl_seconds=cache_ttl_seconds, max_bytes=cache_max_bytes, ) + if provider == "tavily" and not self._api_key: + logger.warning("WebFetchTool: provider='tavily' but api_key is " + "missing; calls will return an error.") @staticmethod def _clean_domains(value: Optional[List[str]]) -> Optional[List[str]]: @@ -551,6 +579,8 @@ async def _fetch_and_build(self, url: str) -> FetchResult: """Issue a streaming GET and map the response to :class:`FetchResult`.""" start = time.monotonic() try: + if self._provider == "tavily": + return await self._extract_with_tavily(url, start) return await self._do_fetch(url, start) except _DomainBlockedError as e: return FetchResult( @@ -579,6 +609,71 @@ async def _fetch_and_build(self, url: str) -> FetchResult: duration_ms=int((time.monotonic() - start) * 1000), ) + async def _extract_with_tavily(self, url: str, start: float) -> FetchResult: + """Extract one public URL with Tavily.""" + if not self._api_key: + return FetchResult( + url=url, + error=("TAVILY_NOT_CONFIGURED: set api_key or " + "TAVILY_API_KEY"), + duration_ms=int((time.monotonic() - start) * 1000), + ) + await self._check_ssrf(url) + + payload: dict[str, Any] = { + "urls": [url], + "extract_depth": self._tavily_extract_depth, + "include_images": False, + } + payload.update(self._tavily_extra_params) + client = self._get_client() + owns_client = self._http_client is None + try: + response = await client.post( + self._base_url, + json=payload, + timeout=self._timeout, + headers={ + "Authorization": f"Bearer {self._api_key}", + "User-Agent": self._user_agent, + }, + ) + response.raise_for_status() + data = response.json() + finally: + if owns_client: + await client.aclose() + + results = data.get("results") or [] + result_data = next( + (item for item in results if isinstance(item, dict)), + None, + ) + if result_data is None: + failed = data.get("failed_results") or data.get("failed_urls") or [] + detail = failed[0] if isinstance(failed, list) and failed else data.get("detail") or data.get("error") + return FetchResult( + url=url, + error=f"TAVILY_EXTRACT_ERROR: {detail or 'no content returned'}", + duration_ms=int((time.monotonic() - start) * 1000), + ) + + final_url = str(result_data.get("url") or url) + content = str(result_data.get("raw_content") or result_data.get("content") or "") + encoded = content.encode("utf-8", errors="ignore") + if self._max_response_bytes > 0 and len(encoded) > self._max_response_bytes: + encoded = encoded[:self._max_response_bytes] + content = encoded.decode("utf-8", errors="ignore") + return FetchResult( + url=final_url, + status_code=200, + status_text="OK", + content_type="text/markdown", + content=content, + bytes=len(content.encode("utf-8", errors="ignore")), + duration_ms=int((time.monotonic() - start) * 1000), + ) + async def _do_fetch(self, url: str, start: float) -> FetchResult: """Drive the request + manual-redirect loop and build the result.""" client = self._get_client() diff --git a/trpc_agent_sdk/tools/_websearch_tool.py b/trpc_agent_sdk/tools/_websearch_tool.py index f40e1101..30731fbb 100644 --- a/trpc_agent_sdk/tools/_websearch_tool.py +++ b/trpc_agent_sdk/tools/_websearch_tool.py @@ -6,8 +6,8 @@ """Web search tool for TRPC Agent framework. Provides a client-side :class:`WebSearchTool` that lets LLMs search the -public web for up-to-date information. Two pluggable provider backends -are supported, ``duckduckgo`` and ``google search``: +public web for up-to-date information. Three pluggable provider backends +are supported, ``duckduckgo``, ``google search``, and ``tavily``: 1. ``duckduckgo`` — DuckDuckGo(DDG) Instant Answer API. Keyless, good for factual/encyclopedic/definition lookups. Returns curated instant @@ -15,6 +15,8 @@ 2. ``google`` — Google Custom Search (CSE). Requires ``api_key`` + ``engine_id``; returns true web results with snippet support, domain filtering and language targeting. +3. ``tavily`` — Tavily Search API. Requires ``api_key``; returns LLM-ready + web results and optionally direct image URLs. """ from __future__ import annotations @@ -64,12 +66,15 @@ _DDG_BASE_URL = "https://api.duckduckgo.com" # Google Custom Search base URL _GOOGLE_BASE_URL = "https://www.googleapis.com/customsearch/v1" +# Tavily Search base URL +_TAVILY_BASE_URL = "https://api.tavily.com/search" # Description shown to the LLM as part of the tool schema. _BASE_DESCRIPTION = """\ Search the public web and use the results to inform responses. - Provides up-to-date information for current events and recent data. - Returns structured results as {title, url, snippet} plus, for DuckDuckGo, an instant-answer summary. All URLs are citable as markdown hyperlinks. +- Tavily can additionally return direct image URLs when `include_images=true`. - Use this tool for accessing information beyond the model's knowledge cutoff. - Each invocation performs a single search request; prefer one well-formed query over many narrow ones. @@ -90,7 +95,7 @@ for the required 'Sources:' format.\ """ -ProviderType = Literal["duckduckgo", "google"] +ProviderType = Literal["duckduckgo", "google", "tavily"] class SearchHit(BaseModel): @@ -101,6 +106,13 @@ class SearchHit(BaseModel): snippet: str = Field(default="", description="Short description / excerpt") +class ImageHit(BaseModel): + """A direct image result returned by a provider that supports images.""" + + url: str = Field(default="", description="Direct image URL") + description: str = Field(default="", description="Provider-supplied image description") + + class WebSearchResult(BaseModel): """Structured output of :class:`WebSearchTool`.""" @@ -112,6 +124,12 @@ class WebSearchResult(BaseModel): summary: str = "" +class TavilyWebSearchResult(WebSearchResult): + """Tavily search output, including optional direct image URLs.""" + + images: List[ImageHit] = Field(default_factory=list) + + def _truncate(text: str, limit: int) -> str: """Truncate text to a maximum length.""" text = (text or "").strip() @@ -259,7 +277,7 @@ class WebSearchTool(BaseTool): """LLM tool that searches the public web. The WebSearchTool enables LLM agents to search the public web using major search engines - such as DuckDuckGo (default, no API key required) and Google Custom Search (API key required). + such as DuckDuckGo (default, no API key required), Google Custom Search, and Tavily. It retrieves up-to-date information including titles, URLs, and content snippets, and also provides instant-answer summaries when available (e.g., via DuckDuckGo). This tool is best used for queries about recent events, new releases, factual lookups, or definitions that benefit @@ -268,8 +286,9 @@ class WebSearchTool(BaseTool): sources as Markdown hyperlinks in the final output. Args: - provider: Backend name, ``"duckduckgo"`` (default) or ``"google"``. - api_key: Google CSE API key; falls back to ``GOOGLE_CSE_API_KEY``. + provider: Backend name: ``"duckduckgo"`` (default), ``"google"``, or ``"tavily"``. + api_key: Provider API key. Falls back to ``GOOGLE_CSE_API_KEY`` for + Google or ``TAVILY_API_KEY`` for Tavily. engine_id: Google CSE engine id (``cx``); falls back to ``GOOGLE_CSE_ENGINE_ID``. results_num: Default result count, clamped to ``[1, _MAX_COUNT]``. snippet_len: Max snippet length, clamped to ``[1, _MAX_SNIPPET_LEN]``. @@ -302,6 +321,7 @@ def __init__( dedup_urls: bool = True, ddg_extra_params: Optional[dict[str, Any]] = None, google_extra_params: Optional[dict[str, Any]] = None, + tavily_extra_params: Optional[dict[str, Any]] = None, filters_name: Optional[List[str]] = None, filters: Optional[List[BaseFilter]] = None, ) -> None: @@ -313,14 +333,26 @@ def __init__( filters=filters, ) + if provider not in ("duckduckgo", "google", "tavily"): + raise ValueError(f"Unsupported web search provider: {provider!r}") self._provider: ProviderType = provider - self._api_key = api_key or os.environ.get("GOOGLE_CSE_API_KEY", "") + if provider == "google": + self._api_key = api_key or os.environ.get("GOOGLE_CSE_API_KEY", "") + elif provider == "tavily": + self._api_key = api_key or os.environ.get("TAVILY_API_KEY", "") + else: + self._api_key = api_key or "" self._engine_id = engine_id or os.environ.get("GOOGLE_CSE_ENGINE_ID", "") self._results_num = max(1, min(int(results_num), _MAX_COUNT)) self._snippet_len = max(1, min(int(snippet_len), _MAX_SNIPPET_LEN)) self._title_len = max(1, min(int(title_len), _MAX_TITLE_LEN)) self._timeout = float(timeout) - self._base_url = base_url or (_GOOGLE_BASE_URL if provider == "google" else _DDG_BASE_URL) + default_base_urls = { + "duckduckgo": _DDG_BASE_URL, + "google": _GOOGLE_BASE_URL, + "tavily": _TAVILY_BASE_URL, + } + self._base_url = base_url or default_base_urls[provider] self._user_agent = user_agent self._proxy = proxy self._lang = lang @@ -328,61 +360,72 @@ def __init__( self._dedup_urls = bool(dedup_urls) self._ddg_extra_params = ddg_extra_params or {} self._google_extra_params = google_extra_params or {} + self._tavily_extra_params = tavily_extra_params or {} if provider == "google" and not (self._api_key and self._engine_id): logger.warning("WebSearchTool: provider='google' but api_key or " "engine_id is missing; calls will return an error.") + if provider == "tavily" and not self._api_key: + logger.warning("WebSearchTool: provider='tavily' but api_key is " + "missing; calls will return an error.") @override def _get_declaration(self) -> FunctionDeclaration: + properties = { + "query": + Schema( + type=Type.STRING, + description=("Required. Search query to send to the provider. Min 2 chars. " + "Include year/version for recent topics. " + "Example: 'Python 3.13 release notes', 'OpenAI GPT-5 pricing 2026', " + "'FastAPI websocket auth', 'vector database definition'."), + ), + "count": + Schema( + type=Type.INTEGER, + description=(f"Optional. Max results to return, 1-{_MAX_COUNT} (clamped). " + f"Default: {self._results_num}. Prefer small values to save context. " + f"Example: 3, 5."), + minimum=1, + maximum=_MAX_COUNT, + ), + "allowed_domains": + Schema( + type=Type.ARRAY, + items=Schema(type=Type.STRING), + description=("Optional. Whitelist of domains, host only (subdomain-aware, " + "'www.' stripped). Mutually exclusive with blocked_domains. " + "Default: None. " + "Example: ['python.org'], ['github.com', 'stackoverflow.com']."), + ), + "blocked_domains": + Schema( + type=Type.ARRAY, + items=Schema(type=Type.STRING), + description=("Optional. Blacklist of domains (same matching as allowed_domains). " + "Mutually exclusive with allowed_domains. Default: None. " + "Example: ['pinterest.com'], ['content-farm.net']."), + ), + "lang": + Schema( + type=Type.STRING, + description=("Optional. Language hint for the provider (Google CSE 'hl'); " + "ignored by DuckDuckGo and Tavily. Default: tool-level lang or unset. " + "Example: 'en', 'zh-CN', 'ja'."), + ), + } + if self._provider == "tavily": + properties["include_images"] = Schema( + type=Type.BOOLEAN, + description=("Optional. Ask Tavily to return direct image URLs. Default: false. " + "Use only when images are useful in the answer."), + ) return FunctionDeclaration( name="websearch", description=_BASE_DESCRIPTION, parameters=Schema( type=Type.OBJECT, - properties={ - "query": - Schema( - type=Type.STRING, - description=("Required. Search query to send to the provider. Min 2 chars. " - "Include year/version for recent topics. " - "Example: 'Python 3.13 release notes', 'OpenAI GPT-5 pricing 2026', " - "'FastAPI websocket auth', 'vector database definition'."), - ), - "count": - Schema( - type=Type.INTEGER, - description=(f"Optional. Max results to return, 1-{_MAX_COUNT} (clamped). " - f"Default: {self._results_num}. Prefer small values to save context. " - f"Example: 3, 5."), - minimum=1, - maximum=_MAX_COUNT, - ), - "allowed_domains": - Schema( - type=Type.ARRAY, - items=Schema(type=Type.STRING), - description=("Optional. Whitelist of domains, host only (subdomain-aware, " - "'www.' stripped). Mutually exclusive with blocked_domains. " - "Default: None. " - "Example: ['python.org'], ['github.com', 'stackoverflow.com']."), - ), - "blocked_domains": - Schema( - type=Type.ARRAY, - items=Schema(type=Type.STRING), - description=("Optional. Blacklist of domains (same matching as allowed_domains). " - "Mutually exclusive with allowed_domains. Default: None. " - "Example: ['pinterest.com'], ['content-farm.net']."), - ), - "lang": - Schema( - type=Type.STRING, - description=("Optional. Language hint for the provider (Google CSE 'hl'); " - "ignored by DuckDuckGo. Default: tool-level lang or unset. " - "Example: 'en', 'zh-CN', 'ja'."), - ), - }, + properties=properties, required=["query"], ), ) @@ -454,10 +497,19 @@ async def _run_async_impl(self, *, tool_context: InvocationContext, args: dict[s n = max(1, min(n, _MAX_COUNT)) lang = (args.get("lang") or self._lang or "").strip() or None + include_images = bool(args.get("include_images", False)) try: if self._provider == "duckduckgo": result = await self._search_duckduckgo(query, n, allowed, blocked) + elif self._provider == "tavily": + result = await self._search_tavily( + query, + n, + allowed, + blocked, + include_images, + ) else: result = await self._search_google(query, n, allowed, blocked, lang) except httpx.HTTPError as e: @@ -497,6 +549,138 @@ async def _get_json(self, url: str, params: dict[str, Any]) -> dict[str, Any]: if close: await client.aclose() + async def _post_json( + self, + url: str, + payload: dict[str, Any], + *, + headers: Optional[dict[str, str]] = None, + ) -> dict[str, Any]: + """Issue a POST and decode the JSON body.""" + client = self._get_client() + close = self._http_client is None + request_headers = {"User-Agent": self._user_agent} + if headers: + request_headers.update(headers) + try: + resp = await client.post( + url, + json=payload, + timeout=self._timeout, + headers=request_headers, + ) + resp.raise_for_status() + return resp.json() + finally: + if close: + await client.aclose() + + async def _search_tavily( + self, + query: str, + n: int, + allowed: Optional[List[str]], + blocked: Optional[List[str]], + include_images: bool, + ) -> TavilyWebSearchResult: + """Hit the Tavily Search API.""" + if not self._api_key: + return TavilyWebSearchResult( + query=query, + provider="tavily", + results=[], + summary=("Tavily provider is not configured: set api_key " + "or TAVILY_API_KEY."), + images=[], + ) + + payload: dict[str, Any] = { + "query": query, + "max_results": n, + "search_depth": "basic", + "include_answer": False, + "include_raw_content": False, + "include_images": include_images, + } + if allowed: + payload["include_domains"] = allowed + if blocked: + payload["exclude_domains"] = blocked + payload.update(self._tavily_extra_params) + + data = await self._post_json( + self._base_url, + payload, + headers={"Authorization": f"Bearer {self._api_key}"}, + ) + if err := (data.get("error") or data.get("detail")): + return TavilyWebSearchResult( + query=query, + provider="tavily", + results=[], + summary=f"Tavily Search API error: {err}", + images=[], + ) + + hits: List[SearchHit] = [] + seen: set[str] = set() + for item in data.get("results") or []: + if not isinstance(item, dict): + continue + url = str(item.get("url") or "").strip() + if _is_blocked(url, allowed, blocked): + continue + if self._dedup_urls: + key = _dedup_key(url) + if key in seen: + continue + seen.add(key) + hits.append( + SearchHit( + title=_truncate(str(item.get("title") or ""), self._title_len), + url=url, + snippet=_truncate(str(item.get("content") or ""), self._snippet_len), + )) + if len(hits) >= n: + break + + images: List[ImageHit] = [] + seen_images: set[str] = set() + if include_images: + image_candidates = list(data.get("images") or []) + for result_item in data.get("results") or []: + if isinstance(result_item, dict): + image_candidates.extend(result_item.get("images") or []) + for item in image_candidates: + if isinstance(item, str): + image_url = item.strip() + description = "" + elif isinstance(item, dict): + image_url = str(item.get("url") or "").strip() + description = str(item.get("description") or "").strip() + else: + continue + if not image_url or image_url in seen_images: + continue + if _extract_domain_from_url(image_url) == "": + continue + seen_images.add(image_url) + images.append(ImageHit( + url=image_url, + description=_truncate(description, self._snippet_len), + )) + if len(images) >= n: + break + + answer = str(data.get("answer") or "").strip() + return TavilyWebSearchResult( + query=query, + provider="tavily", + results=hits, + summary=_truncate(answer, self._snippet_len), + images=images, + ) + async def _search_duckduckgo( self, query: str,