library-desk/src/clients/content_extractor.py

"""
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 '<raw html>'}")
            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 '<raw html>'}: {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")