""" Content extraction client for Library Desk. A reusable Trafilatura wrapper that can be used throughout library-desk: - RAG search service (extract content from search results) - Ingestion service (extract content from URLs) - Standalone endpoint (ad-hoc content extraction) Hardening notes: - Pages are fetched with httpx.AsyncClient under real connect/read timeouts on the event loop. Only the CPU-bound Trafilatura parse runs in the thread pool, so a slow server can no longer pin a worker thread for the duration of a blind blocking download (trafilatura.fetch_url had no caller-side timeout control and kept downloading after asyncio.wait_for gave up). - Trafilatura runs ONCE per document (bare_extraction returns text and metadata together); the old code ran extract() twice (the XML pass was computed and thrown away) plus bare_extraction — three full parses. - extract_batch caps the number of full-page extractions per call; overflow URLs are returned as unsuccessful results so callers fall back to the search-engine snippet. - Thread-pool queue depth is logged so extraction backpressure is visible. """ import asyncio import logging from concurrent.futures import ThreadPoolExecutor from typing import List, Optional import httpx import trafilatura from src.models.content import ContentExtractionResult logger = logging.getLogger(__name__) # Modest but honest identification; some sites reject empty user agents. DEFAULT_HEADERS = { "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) LibraryDesk-ContentExtractor" } # Downloads are streamed and ABORTED past this many bytes (protects memory # and bandwidth during the fetch, and keeps Trafilatura from pinning a # worker thread on a multi-hundred-MB response). MAX_RESPONSE_BYTES = 5 * 1024 * 1024 class ContentExtractor: """ Generic content extraction client using Trafilatura. Fetches pages asynchronously via httpx and runs Trafilatura's synchronous extraction in a thread pool, with support for parallel batch processing and configurable timeouts. """ # Hard cap on full-page extractions per extract_batch call (e.g. one # web-search leg). Overflow URLs get an unsuccessful result and callers # fall back to the search snippet. DEFAULT_MAX_URLS_PER_BATCH = 8 def __init__( self, timeout: int = 5, max_length: int = 2000, max_workers: int = 10, connect_timeout: float = 3.0, max_urls_per_batch: Optional[int] = None ): """ Initialize ContentExtractor. Args: timeout: Per-URL read/parse timeout in seconds max_length: Maximum content length to return (truncated if longer) max_workers: Max concurrent extractions for batch operations connect_timeout: TCP/TLS connect timeout in seconds max_urls_per_batch: Cap on full-page extractions per extract_batch call (None = DEFAULT_MAX_URLS_PER_BATCH) """ self.timeout = timeout self.max_length = max_length self.max_urls_per_batch = ( max_urls_per_batch if max_urls_per_batch is not None else self.DEFAULT_MAX_URLS_PER_BATCH ) self._executor = ThreadPoolExecutor(max_workers=max_workers) self._http = httpx.AsyncClient( timeout=httpx.Timeout(timeout, connect=connect_timeout), follow_redirects=True, headers=DEFAULT_HEADERS, limits=httpx.Limits( max_connections=max_workers, max_keepalive_connections=max_workers ), ) logger.info( f"Initialized ContentExtractor: timeout={timeout}s, " f"connect_timeout={connect_timeout}s, max_length={max_length}, " f"max_workers={max_workers}, max_urls_per_batch={self.max_urls_per_batch}" ) async def _fetch(self, url: str) -> Optional[str]: """ Fetch a URL asynchronously under real connect/read timeouts. The body is STREAMED and the download is aborted as soon as MAX_RESPONSE_BYTES have been received, so the cap bounds memory and bandwidth during the download itself (the old implementation buffered the entire response before truncating, so a multi-hundred-MB URL was still fully downloaded). Returns: Response text (capped at MAX_RESPONSE_BYTES) or None when the response is empty / not OK. Raises: httpx.HTTPError subclasses on timeout/network errors. """ async with self._http.stream("GET", url) as response: if response.status_code != 200: logger.debug( f"Fetch returned status {response.status_code} for {url}" ) return None chunks: List[bytes] = [] received = 0 truncated = False async for chunk in response.aiter_bytes(): if received + len(chunk) >= MAX_RESPONSE_BYTES: chunks.append(chunk[: MAX_RESPONSE_BYTES - received]) truncated = True break chunks.append(chunk) received += len(chunk) body = b"".join(chunks) if not body: return None if truncated: logger.warning( f"Response for {url} exceeds {MAX_RESPONSE_BYTES} bytes; " "download aborted and body truncated" ) # charset comes from the Content-Type header, available before # the body is read. encoding = response.charset_encoding or "utf-8" return body.decode(encoding, errors="replace") def _log_queue_depth(self, context: str) -> None: """Log thread-pool queue depth so extraction backpressure is visible.""" depth = self._executor._work_queue.qsize() if depth > 0: logger.info(f"ContentExtractor thread-pool queue depth ({context}): {depth}") @staticmethod def _extract_html_sync( html: str, url: str, include_metadata: bool, max_length: int ) -> ContentExtractionResult: """ Synchronous Trafilatura pass (runs in the thread pool). Runs bare_extraction ONCE — it returns text and metadata together (the old implementation parsed the document three times). """ try: doc = trafilatura.bare_extraction( html, url=url or None, include_comments=False, include_tables=True, with_metadata=include_metadata ) content = (doc or {}).get("text") or "" if not content: return ContentExtractionResult( url=url, content="", success=False, error="No content extracted" ) # Truncate if needed if len(content) > max_length: content = content[:max_length] + "..." return ContentExtractionResult( url=url, title=doc.get("title") if include_metadata else None, content=content, author=doc.get("author") if include_metadata else None, date=doc.get("date") if include_metadata else None, language=doc.get("language") if include_metadata else None, success=True, error=None ) except Exception as e: logger.error(f"Content extraction failed for {url}: {e}") return ContentExtractionResult( url=url, content="", success=False, error=str(e) ) async def extract( self, url: str, include_metadata: bool = True, max_length: Optional[int] = None ) -> ContentExtractionResult: """ Extract content from a single URL asynchronously. The download happens on the event loop under httpx connect/read timeouts; only the parse occupies a worker thread. Args: url: URL to extract content from include_metadata: Whether to extract title, author, date max_length: Override default max length Returns: ContentExtractionResult with extracted content or error """ try: html = await self._fetch(url) except httpx.TimeoutException: logger.warning(f"Fetch timed out for {url}") return ContentExtractionResult( url=url, content="", success=False, error=f"Fetch timed out after {self.timeout}s" ) except httpx.HTTPError as e: logger.warning(f"Fetch failed for {url}: {e}") return ContentExtractionResult( url=url, content="", success=False, error=f"Failed to fetch URL: {e}" ) if not html: return ContentExtractionResult( url=url, content="", success=False, error="Failed to fetch URL" ) return await self.extract_from_html(html, url, include_metadata, max_length) async def extract_batch( self, urls: List[str], include_metadata: bool = True, max_length: Optional[int] = None, max_urls: Optional[int] = None ) -> List[ContentExtractionResult]: """ Extract content from multiple URLs in parallel. At most max_urls (default: max_urls_per_batch) URLs get a full-page extraction; the rest are returned unsuccessful so callers fall back to their existing snippet. Args: urls: List of URLs to extract content from include_metadata: Whether to extract title, author, date max_length: Override default max length max_urls: Override the per-call full-page extraction cap Returns: List of ContentExtractionResult in same order as input URLs """ cap = max_urls if max_urls is not None else self.max_urls_per_batch fetch_urls = urls[:cap] skipped_urls = urls[cap:] if skipped_urls: logger.info( f"extract_batch capped at {cap} full-page extractions; " f"skipping {len(skipped_urls)} of {len(urls)} URLs" ) self._log_queue_depth("extract_batch") tasks = [ self.extract(url, include_metadata, max_length) for url in fetch_urls ] results = list(await asyncio.gather(*tasks)) results.extend( ContentExtractionResult( url=url, content="", success=False, error=f"Skipped: per-call extraction cap ({cap}) reached" ) for url in skipped_urls ) return results async def extract_from_html( self, html: str, url: str = "", include_metadata: bool = True, max_length: Optional[int] = None ) -> ContentExtractionResult: """ Extract content from raw HTML string. Args: html: Raw HTML content url: Optional URL for reference (not fetched) include_metadata: Whether to extract title, author, date max_length: Override default max length Returns: ContentExtractionResult with extracted content or error """ loop = asyncio.get_event_loop() self._log_queue_depth("extract_from_html") try: return await asyncio.wait_for( loop.run_in_executor( self._executor, self._extract_html_sync, html, url, include_metadata, max_length or self.max_length ), timeout=self.timeout ) except asyncio.TimeoutError: logger.warning(f"Content extraction timed out for {url or ''}") return ContentExtractionResult( url=url, content="", success=False, error=f"Extraction timed out after {self.timeout}s" ) except Exception as e: logger.error(f"Unexpected error extracting {url or ''}: {e}") return ContentExtractionResult( url=url, content="", success=False, error=str(e) ) async def close(self): """Shutdown the HTTP client and the thread pool executor.""" await self._http.aclose() self._executor.shutdown(wait=False) logger.info("ContentExtractor closed")