Add page extraction tool

This commit is contained in:
2026-05-18 22:54:19 +02:00
parent 7147c657ab
commit 68934b23a0
9 changed files with 480 additions and 1 deletions

View File

@@ -8,4 +8,6 @@ BRAVE_SEARCH_COUNTRY=US
BRAVE_SEARCH_LANG=en
BRAVE_SEARCH_UI_LANG=en-US
BRAVE_SEARCH_SAFESEARCH=moderate
PAGE_EXTRACT_MAX_CHARS=12000
PAGE_EXTRACT_RENDER_DYNAMIC=true
# AI_CONFIG={"default_model":"mm","context_messages":12,"prompt":{"root":"prompts/ai","system_file":"system.md","parts_dir":"parts"},"providers":{"minimax":{"api_key":"your_minimax_api_key_here","base_url":"https://api.minimax.io/v1"},"openrouter":{"api_key":"your_openrouter_api_key_here","base_url":"https://openrouter.ai/api/v1"}},"models":{"mm":{"provider":"minimax","model":"MiniMax-Text-01"},"or":{"provider":"openrouter","model":"openai/gpt-4o-mini"}}}

310
app/ai/page_extract.py Normal file
View File

@@ -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_()

View File

@@ -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:

View File

@@ -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] = {

View File

@@ -78,6 +78,16 @@ class AIWebSearchConfig:
return self.enabled and bool(self.api_key)
@dataclass(slots=True)
class AIPageExtractConfig:
enabled: bool = True
timeout: float = 15.0
max_chars: int = 12000
render_dynamic: bool = True
dynamic_wait_ms: int = 1500
block_private_hosts: bool = True
@dataclass(slots=True)
class AIConfig:
default_model: str
@@ -88,6 +98,7 @@ class AIConfig:
workspace_dir: Path = field(default_factory=lambda: Path("workspace"))
tool_max_rounds: int = 10
web_search: AIWebSearchConfig = field(default_factory=AIWebSearchConfig)
page_extract: AIPageExtractConfig = field(default_factory=AIPageExtractConfig)
@dataclass(slots=True)
@@ -231,6 +242,52 @@ def _parse_web_search_config(raw_ai: dict[str, Any]) -> AIWebSearchConfig:
)
def _parse_page_extract_config(raw_ai: dict[str, Any]) -> AIPageExtractConfig:
raw = _parse_mapping(raw_ai.get("page_extract"), field_name="ai.page_extract")
enabled = _parse_bool(raw.get("enabled", True), field_name="ai.page_extract.enabled")
timeout = _parse_float(raw.get("timeout", 15.0), field_name="ai.page_extract.timeout")
if timeout <= 0:
raise ValueError("ai.page_extract.timeout must be > 0")
max_chars = _parse_int(raw.get("max_chars", 12000), field_name="ai.page_extract.max_chars")
if max_chars < 1000:
raise ValueError("ai.page_extract.max_chars must be >= 1000")
render_dynamic = _parse_bool(
raw.get("render_dynamic", True),
field_name="ai.page_extract.render_dynamic",
)
dynamic_wait_ms = _parse_int(
raw.get("dynamic_wait_ms", 1500),
field_name="ai.page_extract.dynamic_wait_ms",
)
dynamic_wait_ms = min(max(dynamic_wait_ms, 0), 10000)
block_private_hosts = _parse_bool(
raw.get("block_private_hosts", True),
field_name="ai.page_extract.block_private_hosts",
)
if os.environ.get("PAGE_EXTRACT_MAX_CHARS") is not None:
max_chars = _parse_int(os.environ["PAGE_EXTRACT_MAX_CHARS"], field_name="PAGE_EXTRACT_MAX_CHARS")
if max_chars < 1000:
raise ValueError("PAGE_EXTRACT_MAX_CHARS must be >= 1000")
if os.environ.get("PAGE_EXTRACT_RENDER_DYNAMIC") is not None:
render_dynamic = _parse_bool(
os.environ["PAGE_EXTRACT_RENDER_DYNAMIC"],
field_name="PAGE_EXTRACT_RENDER_DYNAMIC",
)
return AIPageExtractConfig(
enabled=enabled,
timeout=timeout,
max_chars=max_chars,
render_dynamic=render_dynamic,
dynamic_wait_ms=dynamic_wait_ms,
block_private_hosts=block_private_hosts,
)
def _parse_accounts(raw_accounts: Any, *, default_api_id: int | None, default_api_hash: str | None, default_session_dir: Path) -> list[AccountConfig]:
if raw_accounts is None:
return []
@@ -384,6 +441,7 @@ def _parse_ai_config(raw_ai: Any, *, root: Path) -> AIConfig | None:
raise ValueError("ai.tool_max_rounds must be >= 1")
web_search = _parse_web_search_config(raw_ai)
page_extract = _parse_page_extract_config(raw_ai)
return AIConfig(
default_model=default_model,
@@ -394,6 +452,7 @@ def _parse_ai_config(raw_ai: Any, *, root: Path) -> AIConfig | None:
workspace_dir=workspace_dir,
tool_max_rounds=tool_max_rounds,
web_search=web_search,
page_extract=page_extract,
)

View File

@@ -29,6 +29,14 @@
"timeout": 10.0,
"extra_snippets": false
},
"page_extract": {
"enabled": true,
"timeout": 15.0,
"max_chars": 12000,
"render_dynamic": true,
"dynamic_wait_ms": 1500,
"block_private_hosts": true
},
"prompt": {
"root": "prompts/ai",
"system_file": "system.md",

View File

@@ -3,4 +3,5 @@ tgcrypto
python-dotenv
openai
httpx
playwright
elevenlabs

View File

@@ -29,6 +29,10 @@ def _raw_ai_config() -> dict:
"count": 99,
"safesearch": "strict",
},
"page_extract": {
"max_chars": 5000,
"render_dynamic": False,
},
}
@@ -41,6 +45,8 @@ class ConfigWebSearchTests(unittest.TestCase):
self.assertEqual(config.web_search.api_key, "config-key")
self.assertEqual(config.web_search.count, 20)
self.assertEqual(config.web_search.safesearch, "strict")
self.assertEqual(config.page_extract.max_chars, 5000)
self.assertFalse(config.page_extract.render_dynamic)
def test_web_search_env_overrides(self) -> None:
env = {
@@ -50,6 +56,8 @@ class ConfigWebSearchTests(unittest.TestCase):
"BRAVE_SEARCH_LANG": "pl",
"BRAVE_SEARCH_UI_LANG": "pl-PL",
"BRAVE_SEARCH_SAFESEARCH": "off",
"PAGE_EXTRACT_MAX_CHARS": "9000",
"PAGE_EXTRACT_RENDER_DYNAMIC": "true",
}
with patch.dict("os.environ", env, clear=True):
config = _parse_ai_config(_raw_ai_config(), root=Path.cwd())
@@ -61,6 +69,8 @@ class ConfigWebSearchTests(unittest.TestCase):
self.assertEqual(config.web_search.search_lang, "pl")
self.assertEqual(config.web_search.ui_lang, "pl-PL")
self.assertEqual(config.web_search.safesearch, "off")
self.assertEqual(config.page_extract.max_chars, 9000)
self.assertTrue(config.page_extract.render_dynamic)
if __name__ == "__main__":

View File

@@ -0,0 +1,45 @@
from __future__ import annotations
import unittest
from app.ai.page_extract import extract_readable_text, normalize_url, validate_public_url
class PageExtractTests(unittest.TestCase):
def test_extract_readable_text_skips_scripts_and_styles(self) -> None:
title, text = extract_readable_text(
"""
<html>
<head>
<title>Example Page</title>
<style>.hidden { display: none; }</style>
</head>
<body>
<h1>Hello</h1>
<script>alert('nope')</script>
<p>Readable text.</p>
</body>
</html>
"""
)
self.assertEqual(title, "Example Page")
self.assertIn("Hello", text)
self.assertIn("Readable text.", text)
self.assertNotIn("alert", text)
self.assertNotIn("display", text)
def test_normalize_url_rejects_non_web_schemes(self) -> None:
with self.assertRaisesRegex(ValueError, "http and https"):
normalize_url("file:///etc/passwd")
def test_validate_public_url_blocks_localhost(self) -> None:
with self.assertRaisesRegex(ValueError, "blocked"):
validate_public_url("http://127.0.0.1:8000", block_private_hosts=True)
def test_validate_public_url_can_allow_private_hosts(self) -> None:
validate_public_url("http://127.0.0.1:8000", block_private_hosts=False)
if __name__ == "__main__":
unittest.main()