Oleksandr Kozachuk
237 lines
9.9 KiB
Python
237 lines
9.9 KiB
Python
"""URL reading tool (SPEC-011, FDB-018).
|
|
|
|
The model calls `fetch_url`; this module fetches safely and returns
|
|
readable text plus prominent image URLs. Web pages are hostile input:
|
|
every fetch is SSRF-guarded (no private/loopback/link-local targets,
|
|
http/https only, redirects re-validated) and every byte of text is
|
|
sanitized before it can reach the prompt.
|
|
"""
|
|
|
|
import ipaddress
|
|
import logging
|
|
import re
|
|
import socket
|
|
from html.parser import HTMLParser
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple
|
|
from urllib.parse import urljoin, urlparse
|
|
|
|
import aiohttp
|
|
|
|
from .ai_responder import sanitize_external_text
|
|
from .httpread import read_capped
|
|
|
|
DEFAULT_MAX_BYTES = 2 * 1024 * 1024
|
|
DEFAULT_MAX_CHARS = 8000 # URL-08: budget goes to content now, not chrome
|
|
DEFAULT_MAX_IMAGES = 2
|
|
FETCH_TIMEOUT_S = 15
|
|
MAX_REDIRECTS = 5
|
|
|
|
FETCH_URL_TOOL = {
|
|
"name": "fetch_url",
|
|
"description": "Fetch a public web page and return its readable text plus prominent image links. "
|
|
"Use when the user shares a URL and asks about it, or to get details behind a news link.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"url": {"type": "string", "description": "The http/https URL to read."}},
|
|
"required": ["url"],
|
|
},
|
|
}
|
|
|
|
|
|
_META_REFRESH_URL = re.compile(r"url\s*=\s*['\"]?([^'\";\s]+)", re.I)
|
|
|
|
|
|
_SKIP_TAGS = ("script", "style", "noscript", "svg", "nav", "header", "footer", "aside", "form", "select", "button")
|
|
_BLOCK_TAGS = ("p", "li", "div", "section", "article", "td", "ul", "ol", "table", "h1", "h2", "h3", "h4", "h5", "h6")
|
|
_LINK_DENSITY_MAX = 0.6 # boilerplate: block mostly link text ... (URL-08)
|
|
_LINK_BLOCK_MAX_CHARS = 200 # ... AND short (menus, related lists); long linky paragraphs survive
|
|
|
|
|
|
class _Extractor(HTMLParser):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self._skip = 0
|
|
self._links = 0
|
|
self._buf: List[str] = []
|
|
self._buf_link_chars = 0
|
|
self.blocks: List[Tuple[str, int]] = [] # (text, chars inside <a>)
|
|
self.images: List[str] = []
|
|
self.og_image: Optional[str] = None
|
|
self.refresh_url: Optional[str] = None
|
|
|
|
def _flush(self) -> None:
|
|
text = " ".join(self._buf).strip()
|
|
if text:
|
|
self.blocks.append((text, self._buf_link_chars))
|
|
self._buf, self._buf_link_chars = [], 0
|
|
|
|
def handle_starttag(self, tag: str, attrs) -> None:
|
|
if tag in _SKIP_TAGS:
|
|
self._skip += 1
|
|
if tag == "a":
|
|
self._links += 1
|
|
if tag in _BLOCK_TAGS:
|
|
self._flush()
|
|
attr = dict(attrs)
|
|
src = attr.get("src")
|
|
if tag == "img" and src:
|
|
self.images.append(src)
|
|
if tag == "meta" and attr.get("property") == "og:image" and attr.get("content"):
|
|
self.og_image = attr["content"]
|
|
# meta-refresh redirect (link shorteners, getnews stubs) — URL-04
|
|
content = attr.get("content")
|
|
if tag == "meta" and (attr.get("http-equiv") or "").lower() == "refresh" and content:
|
|
match = _META_REFRESH_URL.search(content)
|
|
if match and self.refresh_url is None:
|
|
self.refresh_url = match.group(1)
|
|
|
|
def handle_endtag(self, tag: str) -> None:
|
|
if tag in _SKIP_TAGS and self._skip > 0:
|
|
self._skip -= 1
|
|
if tag == "a" and self._links > 0:
|
|
self._links -= 1
|
|
if tag in _BLOCK_TAGS:
|
|
self._flush()
|
|
|
|
def handle_data(self, data: str) -> None:
|
|
if self._skip == 0 and data.strip():
|
|
self._buf.append(data.strip())
|
|
if self._links > 0:
|
|
self._buf_link_chars += len(data.strip())
|
|
|
|
def content_parts(self) -> List[str]:
|
|
"""Blocks minus boilerplate: short blocks dominated by link text are chrome (URL-08)."""
|
|
self._flush()
|
|
out = []
|
|
for text, link_chars in self.blocks:
|
|
if link_chars / max(1, len(text)) > _LINK_DENSITY_MAX and len(text) < _LINK_BLOCK_MAX_CHARS:
|
|
continue
|
|
out.append(text)
|
|
return out
|
|
|
|
|
|
def _ip_is_public(ip_str: str) -> bool:
|
|
try:
|
|
ip = ipaddress.ip_address(ip_str)
|
|
except ValueError:
|
|
return False
|
|
return not (ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_multicast or ip.is_reserved or ip.is_unspecified)
|
|
|
|
|
|
def guard_url(url: str) -> Optional[str]:
|
|
"""Return None if safe to fetch, else a human-readable refusal reason (URL-02/03)."""
|
|
parsed = urlparse(url)
|
|
if parsed.scheme not in ("http", "https"):
|
|
return f"refused scheme {parsed.scheme!r} (only http/https)"
|
|
host = parsed.hostname
|
|
if not host:
|
|
return "refused: no host"
|
|
try:
|
|
literal = ipaddress.ip_address(host)
|
|
return None if _ip_is_public(str(literal)) else f"refused non-public address {host}"
|
|
except ValueError:
|
|
pass
|
|
try:
|
|
infos = socket.getaddrinfo(host, None)
|
|
except socket.gaierror:
|
|
return f"refused: cannot resolve {host}"
|
|
for info in infos:
|
|
if not _ip_is_public(str(info[4][0])):
|
|
return f"refused: {host} resolves to non-public address"
|
|
return None
|
|
|
|
|
|
class URLReader:
|
|
def __init__(self, config_getter: Callable[[], Dict[str, Any]], image_cache) -> None:
|
|
self._config = config_getter
|
|
self.image_cache = image_cache
|
|
|
|
def enabled(self) -> bool:
|
|
return bool(self._config().get("enable-url-reading", False))
|
|
|
|
async def _get(self, session, url: str, max_bytes: int) -> Tuple[str, bytes, str]:
|
|
"""Manual redirect handling so every hop is re-guarded (URL-04)."""
|
|
current = url
|
|
for _ in range(MAX_REDIRECTS):
|
|
reason = guard_url(current)
|
|
if reason:
|
|
raise ValueError(reason)
|
|
async with session.get(current, allow_redirects=False) as response:
|
|
if response.status in (301, 302, 303, 307, 308) and response.headers.get("Location"):
|
|
current = urljoin(current, response.headers["Location"])
|
|
continue
|
|
response.raise_for_status()
|
|
content_type = str(response.headers.get("Content-Type", "")).split(";")[0].strip().lower()
|
|
return str(response.url), await read_capped(response, max_bytes), content_type
|
|
raise ValueError("too many redirects")
|
|
|
|
async def fetch(self, url: str, channel: str, user: str) -> Dict[str, Any]:
|
|
config = self._config()
|
|
max_bytes = int(config.get("url-max-bytes", DEFAULT_MAX_BYTES))
|
|
timeout = aiohttp.ClientTimeout(total=FETCH_TIMEOUT_S)
|
|
try:
|
|
async with aiohttp.ClientSession(timeout=timeout, headers={"User-Agent": "FjerkroaBot/1.0"}) as session:
|
|
final_url, body, content_type = await self._get(session, url, max_bytes)
|
|
# follow a meta-refresh redirect (link shorteners / getnews stubs), re-guarded — URL-04
|
|
for _ in range(2):
|
|
if content_type.startswith("image/"):
|
|
break
|
|
extractor = self._extract(body.decode("utf-8", "ignore"))
|
|
if not extractor.refresh_url:
|
|
break
|
|
target = urljoin(final_url, extractor.refresh_url)
|
|
if guard_url(target) is not None or target == final_url:
|
|
break
|
|
logging.info(f"url reader: following meta-refresh -> {target}")
|
|
final_url, body, content_type = await self._get(session, target, max_bytes)
|
|
except Exception as err:
|
|
return {"error": str(err)}
|
|
# a URL that IS an image: cache it and hand it over as sight (URL-09)
|
|
if content_type.startswith("image/"):
|
|
vision = []
|
|
if self.image_cache is not None and len(body) < max_bytes: # >= cap means possibly truncated
|
|
sha = self.image_cache.ingest_bytes(body, channel, user, None)
|
|
data_url = self._cached_data_url(sha, channel) if sha else None
|
|
vision = [data_url] if data_url else []
|
|
return {"url": final_url, "text": "(image)", "images_cached": len(vision), "vision": vision}
|
|
html = body.decode("utf-8", "ignore")
|
|
clean = sanitize_external_text(self._to_text(html), int(config.get("url-max-chars", DEFAULT_MAX_CHARS)))
|
|
vision = await self._ingest_images(html, final_url, channel, user)
|
|
return {"url": final_url, "text": clean, "images_cached": len(vision), "vision": vision}
|
|
|
|
def _extract(self, html: str) -> "_Extractor":
|
|
extractor = _Extractor()
|
|
try:
|
|
extractor.feed(html)
|
|
except Exception as err:
|
|
logging.debug(f"html parse failed: {err!r}")
|
|
return extractor
|
|
|
|
def _to_text(self, html: str) -> str:
|
|
return re.sub(r"\s+\n", "\n", " ".join(self._extract(html).content_parts()))
|
|
|
|
async def _ingest_images(self, html: str, base_url: str, channel: str, user: str) -> List[str]:
|
|
"""Cache page images and return their data URLs for vision input (URL-09)."""
|
|
if self.image_cache is None:
|
|
return []
|
|
extractor = self._extract(html)
|
|
candidates = ([extractor.og_image] if extractor.og_image else []) + extractor.images
|
|
limit = int(self._config().get("url-max-images", DEFAULT_MAX_IMAGES))
|
|
data_urls: List[str] = []
|
|
for src in candidates:
|
|
if len(data_urls) >= limit:
|
|
break
|
|
absolute = urljoin(base_url, src)
|
|
if guard_url(absolute) is not None:
|
|
continue
|
|
sha = await self.image_cache.ingest_url(absolute, channel, user, None)
|
|
data_url = self._cached_data_url(sha, channel) if sha else None
|
|
if data_url:
|
|
data_urls.append(data_url)
|
|
return data_urls
|
|
|
|
def _cached_data_url(self, sha: str, channel: str) -> Optional[str]:
|
|
recent = self.image_cache.recent(channel, 8)
|
|
ext = next((row["ext"] for row in recent if row["sha256"] == sha), None)
|
|
return self.image_cache.data_url(sha, ext) if ext else None
|