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)
"""

import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional

import trafilatura

from src.models.content import ContentExtractionResult

logger = logging.getLogger(__name__)


class ContentExtractor:
    """
    Generic content extraction client using Trafilatura.

    Provides async wrappers around Trafilatura's synchronous extraction,
    with support for parallel batch processing and configurable timeouts.
    """

    def __init__(
        self,
        timeout: int = 5,
        max_length: int = 2000,
        max_workers: int = 10
    ):
        """
        Initialize ContentExtractor.

        Args:
            timeout: Per-URL timeout in seconds
            max_length: Maximum content length to return (truncated if longer)
            max_workers: Max concurrent extractions for batch operations
        """
        self.timeout = timeout
        self.max_length = max_length
        self._executor = ThreadPoolExecutor(max_workers=max_workers)
        logger.info(
            f"Initialized ContentExtractor: timeout={timeout}s, "
            f"max_length={max_length}, max_workers={max_workers}"
        )

    def _extract_sync(
        self,
        url: str,
        include_metadata: bool = True,
        max_length: Optional[int] = None
    ) -> ContentExtractionResult:
        """
        Synchronous extraction (runs in thread pool).

        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
        """
        effective_max_length = max_length or self.max_length

        try:
            # Fetch the URL
            downloaded = trafilatura.fetch_url(url)
            if not downloaded:
                return ContentExtractionResult(
                    url=url,
                    content="",
                    success=False,
                    error="Failed to fetch URL"
                )

            # Extract content
            content = trafilatura.extract(
                downloaded,
                include_comments=False,
                include_tables=True,
                output_format='txt'
            )

            if not content:
                return ContentExtractionResult(
                    url=url,
                    content="",
                    success=False,
                    error="No content extracted"
                )

            # Truncate if needed
            if len(content) > effective_max_length:
                content = content[:effective_max_length] + "..."

            # Extract metadata if requested
            title = None
            author = None
            date = None
            language = None

            if include_metadata:
                metadata = trafilatura.extract(
                    downloaded,
                    output_format='xml',
                    include_comments=False
                )
                # Parse metadata from XML if available
                # trafilatura.extract with output_format='xml' returns XML with metadata
                # For simplicity, we'll use bare_extraction which returns a dict
                try:
                    meta_dict = trafilatura.bare_extraction(
                        downloaded,
                        include_comments=False
                    )
                    if meta_dict:
                        title = meta_dict.get('title')
                        author = meta_dict.get('author')
                        date = meta_dict.get('date')
                        language = meta_dict.get('language')
                except Exception as e:
                    logger.debug(f"Metadata extraction failed for {url}: {e}")

            return ContentExtractionResult(
                url=url,
                title=title,
                content=content,
                author=author,
                date=date,
                language=language,
                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.

        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
        """
        loop = asyncio.get_event_loop()

        try:
            result = await asyncio.wait_for(
                loop.run_in_executor(
                    self._executor,
                    self._extract_sync,
                    url,
                    include_metadata,
                    max_length
                ),
                timeout=self.timeout
            )
            return result
        except asyncio.TimeoutError:
            logger.warning(f"Content extraction timed out for {url}")
            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}: {e}")
            return ContentExtractionResult(
                url=url,
                content="",
                success=False,
                error=str(e)
            )

    async def extract_batch(
        self,
        urls: List[str],
        include_metadata: bool = True,
        max_length: Optional[int] = None
    ) -> List[ContentExtractionResult]:
        """
        Extract content from multiple URLs in parallel.

        Args:
            urls: List of URLs to extract content from
            include_metadata: Whether to extract title, author, date
            max_length: Override default max length

        Returns:
            List of ContentExtractionResult in same order as input URLs
        """
        tasks = [
            self.extract(url, include_metadata, max_length)
            for url in urls
        ]
        results = await asyncio.gather(*tasks)
        return list(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
        """
        effective_max_length = max_length or self.max_length

        def _extract():
            try:
                content = trafilatura.extract(
                    html,
                    include_comments=False,
                    include_tables=True,
                    output_format='txt'
                )

                if not content:
                    return ContentExtractionResult(
                        url=url,
                        content="",
                        success=False,
                        error="No content extracted from HTML"
                    )

                # Truncate if needed
                if len(content) > effective_max_length:
                    content = content[:effective_max_length] + "..."

                # Extract metadata
                title = None
                author = None
                date = None
                language = None

                if include_metadata:
                    try:
                        meta_dict = trafilatura.bare_extraction(
                            html,
                            include_comments=False
                        )
                        if meta_dict:
                            title = meta_dict.get('title')
                            author = meta_dict.get('author')
                            date = meta_dict.get('date')
                            language = meta_dict.get('language')
                    except Exception as e:
                        logger.debug(f"Metadata extraction failed: {e}")

                return ContentExtractionResult(
                    url=url,
                    title=title,
                    content=content,
                    author=author,
                    date=date,
                    language=language,
                    success=True,
                    error=None
                )

            except Exception as e:
                logger.error(f"HTML content extraction failed: {e}")
                return ContentExtractionResult(
                    url=url,
                    content="",
                    success=False,
                    error=str(e)
                )

        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(self._executor, _extract)

    async def close(self):
        """Shutdown the thread pool executor."""
        self._executor.shutdown(wait=False)
        logger.info("ContentExtractor closed")