Add page extraction tool
This commit is contained in:
@@ -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
310
app/ai/page_extract.py
Normal 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_()
|
||||
@@ -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:
|
||||
|
||||
@@ -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] = {
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -3,4 +3,5 @@ tgcrypto
|
||||
python-dotenv
|
||||
openai
|
||||
httpx
|
||||
playwright
|
||||
elevenlabs
|
||||
|
||||
@@ -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__":
|
||||
|
||||
45
tests/test_page_extract.py
Normal file
45
tests/test_page_extract.py
Normal 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()
|
||||
Reference in New Issue
Block a user