diff --git a/src/stirrup/tools/web.py b/src/stirrup/tools/web.py index 2544b54..836182a 100644 --- a/src/stirrup/tools/web.py +++ b/src/stirrup/tools/web.py @@ -170,6 +170,32 @@ async def _resolve_public_addresses(url: httpx.URL) -> tuple[str, ...]: return (str(literal_address),) +HTML_MEDIA_TYPES = frozenset({"text/html", "application/xhtml+xml"}) +TEXT_MEDIA_TYPES = frozenset( + { + "application/json", + "application/xml", + "application/csv", + "application/javascript", + "application/x-ndjson", + "application/x-yaml", + "application/yaml", + } +) + + +def _media_type(response: httpx.Response) -> str: + """Return the lowercased media type of ``response`` without its parameters.""" + return response.headers.get("content-type", "").split(";", 1)[0].strip().lower() + + +def _is_text_media_type(media_type: str) -> bool: + """Whether a response of ``media_type`` is text to return as is rather than HTML to extract.""" + if media_type in HTML_MEDIA_TYPES: + return False + return media_type.startswith("text/") or media_type in TEXT_MEDIA_TYPES or media_type.endswith(("+json", "+xml")) + + def _get_fetch_web_page_tool( client: httpx.AsyncClient, *, @@ -239,7 +265,7 @@ async def _request_hop( raise RuntimeError("Web fetch retry loop completed without a response.") - async def _fetch_with_redirects(url: str) -> bytes: + async def _fetch_with_redirects(url: str) -> httpx.Response: current_url = httpx.URL(url).copy_with(fragment=None) redirects_followed = 0 @@ -247,7 +273,7 @@ async def _fetch_with_redirects(url: str) -> bytes: public_addresses = await _resolve_public_addresses(current_url) response = await _request_hop(current_url, public_addresses) if not response.has_redirect_location: - return response.content + return response if redirects_followed >= WEB_FETCH_MAX_REDIRECTS: raise _WebFetchError(f"Too many redirects (maximum {WEB_FETCH_MAX_REDIRECTS}).") @@ -259,14 +285,21 @@ async def fetch_web_page_executor(params: FetchWebPageParams) -> ToolResult[WebF """Fetch web page and extract main content as markdown using trafilatura.""" try: with fail_after(timeout): - response_body = await _fetch_with_redirects(params.url) - body_md = ( - await to_thread.run_sync( - partial(trafilatura.extract, response_body, output_format="markdown"), - abandon_on_cancel=True, + response = await _fetch_with_redirects(params.url) + media_type = _media_type(response) + if _is_text_media_type(media_type): + # trafilatura only extracts HTML; it returns None for JSON, CSV, XML or plain text. + body_md = response.text + elif media_type in HTML_MEDIA_TYPES or not media_type: + body_md = ( + await to_thread.run_sync( + partial(trafilatura.extract, response.content, output_format="markdown"), + abandon_on_cancel=True, + ) + or "" ) - or "" - ) + else: + raise _WebFetchError(f"Unsupported content type: {media_type}") return ToolResult( content=f"{params.url}" f"{truncate_msg(body_md, MAX_LENGTH_WEB_FETCH_HTML)}", @@ -287,8 +320,9 @@ async def fetch_web_page_executor(params: FetchWebPageParams) -> ToolResult[WebF return Tool[FetchWebPageParams, WebFetchMetadata]( name="fetch_web_page", description=( - "Fetch and extract the main content from a web page as markdown. Cannot reach private, loopback, " - "or other non-public network addresses. Returns body text or error as XML." + "Fetch and extract the main content from a web page as markdown. Text responses such as JSON, CSV, " + "XML or plain text are returned as is. Cannot reach private, loopback, or other non-public network " + "addresses. Returns body text or error as XML." ), parameters=FetchWebPageParams, executor=fetch_web_page_executor, diff --git a/tests/test_web_fetch.py b/tests/test_web_fetch.py index 12e249f..5168509 100644 --- a/tests/test_web_fetch.py +++ b/tests/test_web_fetch.py @@ -67,7 +67,7 @@ def html_response( return httpx.Response( status_code, request=request, - headers=headers, + headers={"content-type": "text/html; charset=utf-8", **(headers or {})}, text="

Example body

", ) @@ -536,3 +536,59 @@ async def handler(request: httpx.Request) -> httpx.Response: assert result.success is False assert len(requests) == web.WEB_FETCH_MAX_REDIRECTS + 1 + + +@pytest.mark.parametrize( + ("content_type", "body"), + [ + ("application/json", '[{"page": 1}, [{"value": 9123.4}]]'), + ("application/vnd.api+json; charset=utf-8", '{"data": []}'), + ("text/csv", "Entity,Year,Life expectancy\nBrazil,2023,76.2\n"), + ("text/plain; charset=utf-8", "Internet Engineering Task Force (IETF)"), + ("application/xml", "1"), + ], +) +async def test_web_fetch_returns_text_responses_as_is( + monkeypatch: MonkeyPatch, + content_type: str, + body: str, +) -> None: + install_dns(monkeypatch, PUBLIC_IPV4) + + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, request=request, headers={"content-type": content_type}, text=body) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + result = await run_web_fetch_tool(_get_fetch_web_page_tool(client), "https://example.com/data") + + assert result.success is True + assert f"{body}" in result.content + + +async def test_web_fetch_extracts_html_responses(monkeypatch: MonkeyPatch) -> None: + install_dns(monkeypatch, PUBLIC_IPV4) + + async def handler(request: httpx.Request) -> httpx.Response: + return html_response(request) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + result = await run_web_fetch_tool(_get_fetch_web_page_tool(client), "https://example.com/page") + + assert result.success is True + assert "Example body" in result.content + assert "
" not in result.content + + +async def test_web_fetch_reports_unsupported_content_types(monkeypatch: MonkeyPatch) -> None: + install_dns(monkeypatch, PUBLIC_IPV4) + + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, request=request, headers={"content-type": "application/pdf"}, content=b"%PDF-1.7 binary" + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + result = await run_web_fetch_tool(_get_fetch_web_page_tool(client), "https://example.com/report.pdf") + + assert result.success is False + assert "Unsupported content type: application/pdf" in result.content