Add page extraction tool

app/ai/page_extract.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from html.parser import HTMLParser
+import ipaddress
+import re
+import socket
+from typing import Any
+from urllib.parse import urljoin, urlparse
+
+import httpx
+
+from app.config import AIPageExtractConfig
+
+TEXT_BLOCK_TAGS = {
+    "address",
+    "article",
+    "aside",
+    "blockquote",
+    "br",
+    "dd",
+    "div",
+    "dl",
+    "dt",
+    "figcaption",
+    "footer",
+    "h1",
+    "h2",
+    "h3",
+    "h4",
+    "h5",
+    "h6",
+    "header",
+    "li",
+    "main",
+    "nav",
+    "ol",
+    "p",
+    "pre",
+    "section",
+    "table",
+    "td",
+    "th",
+    "tr",
+    "ul",
+}
+SKIP_TAGS = {"script", "style", "noscript", "template", "svg", "canvas"}
+
+
+@dataclass(slots=True)
+class PageExtractResult:
+    url: str
+    final_url: str
+    title: str
+    text: str
+    method: str
+    truncated: bool = False
+
+
+class ReadableTextParser(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__(convert_charrefs=True)
+        self.title = ""
+        self._in_title = False
+        self._skip_depth = 0
+        self._parts: list[str] = []
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        tag = tag.lower()
+        if tag in SKIP_TAGS:
+            self._skip_depth += 1
+            return
+        if tag == "title":
+            self._in_title = True
+            return
+        if tag in TEXT_BLOCK_TAGS:
+            self._parts.append("\n")
+
+    def handle_endtag(self, tag: str) -> None:
+        tag = tag.lower()
+        if tag in SKIP_TAGS and self._skip_depth:
+            self._skip_depth -= 1
+            return
+        if tag == "title":
+            self._in_title = False
+            return
+        if tag in TEXT_BLOCK_TAGS:
+            self._parts.append("\n")
+
+    def handle_data(self, data: str) -> None:
+        if self._skip_depth:
+            return
+        text = data.strip()
+        if not text:
+            return
+        if self._in_title:
+            self.title = " ".join((self.title, text)).strip()
+            return
+        self._parts.append(text)
+        self._parts.append(" ")
+
+    def get_text(self) -> str:
+        text = "".join(self._parts)
+        text = re.sub(r"[ \t\r\f\v]+", " ", text)
+        text = re.sub(r"\n\s+", "\n", text)
+        text = re.sub(r"\n{3,}", "\n\n", text)
+        return text.strip()
+
+
+def extract_readable_text(html: str) -> tuple[str, str]:
+    parser = ReadableTextParser()
+    parser.feed(html)
+    parser.close()
+    return parser.title.strip(), parser.get_text()
+
+
+def normalize_url(raw_url: str) -> str:
+    url = str(raw_url or "").strip()
+    if not url:
+        raise ValueError("extract_page: 'url' is required")
+    parsed = urlparse(url)
+    if parsed.scheme not in {"http", "https"}:
+        raise ValueError("extract_page: only http and https URLs are allowed")
+    if not parsed.hostname:
+        raise ValueError("extract_page: URL must include a hostname")
+    return url
+
+
+def validate_public_url(url: str, *, block_private_hosts: bool) -> None:
+    if not block_private_hosts:
+        return
+    parsed = urlparse(url)
+    host = parsed.hostname
+    if not host:
+        raise ValueError("extract_page: URL must include a hostname")
+
+    try:
+        addresses = socket.getaddrinfo(host, parsed.port or 443, type=socket.SOCK_STREAM)
+    except socket.gaierror as exc:
+        raise ValueError(f"extract_page: could not resolve host {host!r}") from exc
+
+    for address in addresses:
+        ip = ipaddress.ip_address(address[4][0])
+        if (
+            ip.is_private
+            or ip.is_loopback
+            or ip.is_link_local
+            or ip.is_multicast
+            or ip.is_reserved
+            or ip.is_unspecified
+        ):
+            raise ValueError("extract_page: private, local, and reserved hosts are blocked")
+
+
+def _truncate(text: str, max_chars: int) -> tuple[str, bool]:
+    if len(text) <= max_chars:
+        return text, False
+    return text[:max_chars].rstrip() + "\n...(truncated)", True
+
+
+def _format_result(result: PageExtractResult) -> str:
+    lines = [
+        f"URL: {result.final_url or result.url}",
+        f"Method: {result.method}",
+    ]
+    if result.title:
+        lines.append(f"Title: {result.title}")
+    if result.truncated:
+        lines.append("Note: content was truncated.")
+    lines.append("")
+    lines.append(result.text or "(no readable text extracted)")
+    return "\n".join(lines)
+
+
+class PageExtractor:
+    def __init__(self, config: AIPageExtractConfig) -> None:
+        self.config = config
+
+    @property
+    def available(self) -> bool:
+        return self.config.enabled
+
+    async def extract(self, arguments: dict[str, Any]) -> PageExtractResult:
+        if not self.config.enabled:
+            raise ValueError("Page extraction is disabled")
+
+        url = normalize_url(str(arguments.get("url") or ""))
+        validate_public_url(url, block_private_hosts=self.config.block_private_hosts)
+
+        render = str(arguments.get("render", "auto")).strip().lower()
+        if render not in {"auto", "never", "always"}:
+            raise ValueError("extract_page: 'render' must be auto, never, or always")
+
+        raw_max_chars = arguments.get("max_chars", self.config.max_chars)
+        try:
+            max_chars = int(raw_max_chars)
+        except (TypeError, ValueError):
+            max_chars = self.config.max_chars
+        max_chars = min(max(max_chars, 1000), 50000)
+
+        if render == "always":
+            return await self._extract_dynamic(url, max_chars=max_chars)
+
+        static_result = await self._extract_static(url, max_chars=max_chars)
+        if render == "never" or not self.config.render_dynamic:
+            return static_result
+        if len(static_result.text) >= 700:
+            return static_result
+
+        try:
+            return await self._extract_dynamic(url, max_chars=max_chars)
+        except ValueError:
+            return static_result
+
+    async def extract_text(self, arguments: dict[str, Any]) -> str:
+        return _format_result(await self.extract(arguments))
+
+    async def _extract_static(self, url: str, *, max_chars: int) -> PageExtractResult:
+        headers = {
+            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
+            "User-Agent": (
+                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+                "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120 Safari/537.36"
+            ),
+        }
+        try:
+            async with httpx.AsyncClient(
+                timeout=self.config.timeout,
+                headers=headers,
+            ) as client:
+                current_url = url
+                for _ in range(6):
+                    validate_public_url(
+                        current_url,
+                        block_private_hosts=self.config.block_private_hosts,
+                    )
+                    response = await client.get(current_url)
+                    if response.status_code not in {301, 302, 303, 307, 308}:
+                        break
+                    location = response.headers.get("location")
+                    if not location:
+                        break
+                    current_url = urljoin(str(response.url), location)
+                else:
+                    raise ValueError("extract_page: too many redirects")
+        except httpx.HTTPError as exc:
+            raise ValueError(f"extract_page: request failed: {exc}") from exc
+
+        if response.status_code >= 400:
+            raise ValueError(f"extract_page: HTTP {response.status_code}")
+
+        content_type = response.headers.get("content-type", "").lower()
+        if "text/html" not in content_type and "application/xhtml" not in content_type:
+            raise ValueError(f"extract_page: unsupported content type: {content_type or 'unknown'}")
+
+        title, text = extract_readable_text(response.text)
+        text, truncated = _truncate(text, max_chars)
+        return PageExtractResult(
+            url=url,
+            final_url=str(response.url),
+            title=title,
+            text=text,
+            method="http",
+            truncated=truncated,
+        )
+
+    async def _extract_dynamic(self, url: str, *, max_chars: int) -> PageExtractResult:
+        try:
+            from playwright.async_api import async_playwright
+        except ModuleNotFoundError as exc:
+            raise ValueError(
+                "extract_page: dynamic rendering requires playwright and browser install"
+            ) from exc
+
+        try:
+            async with async_playwright() as playwright:
+                browser = await playwright.chromium.launch(headless=True)
+                try:
+                    page = await browser.new_page()
+                    if self.config.block_private_hosts:
+                        await page.route("**/*", self._route_public_requests_only)
+                    await page.goto(url, wait_until="networkidle", timeout=int(self.config.timeout * 1000))
+                    if self.config.dynamic_wait_ms:
+                        await page.wait_for_timeout(self.config.dynamic_wait_ms)
+                    title = (await page.title()).strip()
+                    text = await page.locator("body").inner_text(timeout=5000)
+                    final_url = page.url
+                finally:
+                    await browser.close()
+        except Exception as exc:
+            raise ValueError(f"extract_page: dynamic render failed: {exc}") from exc
+
+        text = re.sub(r"\n{3,}", "\n\n", text).strip()
+        text, truncated = _truncate(text, max_chars)
+        return PageExtractResult(
+            url=url,
+            final_url=final_url,
+            title=title,
+            text=text,
+            method="browser",
+            truncated=truncated,
+        )
+
+    async def _route_public_requests_only(self, route: Any) -> None:
+        try:
+            validate_public_url(route.request.url, block_private_hosts=True)
+        except ValueError:
+            await route.abort()
+            return
+        await route.continue_()

app/ai/service.py

@@ -14,7 +14,14 @@ from openai import AsyncOpenAI
 
 from app.ai.parser import AIInvocation
 from app.ai.prompts import PromptBook
-from app.ai.tools import FILE_TOOL_DISPATCH, FILE_TOOL_SCHEMAS, PYTHON_EXEC_SCHEMA, WEB_SEARCH_SCHEMA
+from app.ai.page_extract import PageExtractor
+from app.ai.tools import (
+    FILE_TOOL_DISPATCH,
+    FILE_TOOL_SCHEMAS,
+    PAGE_EXTRACT_SCHEMA,
+    PYTHON_EXEC_SCHEMA,
+    WEB_SEARCH_SCHEMA,
+)
 from app.ai.web_search import BraveWebSearchClient, format_results
 from app.config import AIConfig, AIModelConfig, AIProviderConfig
 
@@ -64,6 +71,7 @@ class AIService:
         self._clients: dict[str, AsyncOpenAI] = {}
         self._system_prompt: str | None = None
         self.web_search = BraveWebSearchClient(config.web_search)
+        self.page_extractor = PageExtractor(config.page_extract)
 
     async def initialize(self) -> None:
         return None
@@ -215,6 +223,8 @@ class AIService:
         tools = FILE_TOOL_SCHEMAS + [PYTHON_EXEC_SCHEMA]
         if self.web_search.available:
             tools.append(WEB_SEARCH_SCHEMA)
+        if self.page_extractor.available:
+            tools.append(PAGE_EXTRACT_SCHEMA)
         return tools
 
     def _get_allowed_roots(self) -> tuple[Path, ...]:
@@ -347,6 +357,8 @@ class AIService:
             elif name == "web_search":
                 response = await self.web_search.search(arguments)
                 result = format_results(response, include_heading=True)
+            elif name == "extract_page":
+                result = await self.page_extractor.extract_text(arguments)
             elif name == "send_file":
                 result = await self._execute_send_file(arguments, ctx, allowed_roots)
             elif name in FILE_TOOL_DISPATCH:

app/ai/tools.py

@@ -273,6 +273,38 @@ WEB_SEARCH_SCHEMA: dict[str, Any] = {
     },
 }
 
+PAGE_EXTRACT_SCHEMA: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "extract_page",
+        "description": (
+            "Extract readable text from a web page URL. Use render='always' for dynamic "
+            "JavaScript-heavy pages, render='never' for simple static pages, or render='auto' by default."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "url": {
+                    "type": "string",
+                    "description": "HTTP or HTTPS page URL to read.",
+                },
+                "render": {
+                    "type": "string",
+                    "enum": ["auto", "never", "always"],
+                    "description": "Whether to use browser rendering for dynamic pages.",
+                },
+                "max_chars": {
+                    "type": "integer",
+                    "description": "Maximum extracted characters to return, 1000-50000.",
+                    "minimum": 1000,
+                    "maximum": 50000,
+                },
+            },
+            "required": ["url"],
+        },
+    },
+}
+
 ToolExecutor = Callable[[dict, tuple[Path, ...]], Awaitable[str]]
 
 FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {