corp-extractor 0.4.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

statement_extractor/document/loader.py

@@ -0,0 +1,303 @@
+"""
+URL loader for fetching and parsing web content.
+
+Orchestrates scraper and PDF parser plugins to load documents from URLs.
+"""
+
+import asyncio
+import logging
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from ..models.document import Document
+from ..pipeline.registry import PluginRegistry
+from ..plugins.base import (
+    BaseScraperPlugin,
+    BasePDFParserPlugin,
+    ContentType,
+    ScraperResult,
+)
+from .html_extractor import extract_text_from_html, extract_article_content
+
+logger = logging.getLogger(__name__)
+
+
+class URLLoaderConfig(BaseModel):
+    """Configuration for URL loading."""
+
+    timeout: float = Field(
+        default=30.0,
+        description="Request timeout in seconds"
+    )
+    use_ocr: bool = Field(
+        default=False,
+        description="Force OCR for PDF parsing"
+    )
+    max_pdf_pages: int = Field(
+        default=500,
+        description="Maximum pages to extract from PDFs"
+    )
+    scraper_plugin: Optional[str] = Field(
+        default=None,
+        description="Specific scraper plugin to use (None = auto-select)"
+    )
+    pdf_parser_plugin: Optional[str] = Field(
+        default=None,
+        description="Specific PDF parser plugin to use (None = auto-select)"
+    )
+    extract_metadata: bool = Field(
+        default=True,
+        description="Extract article metadata from HTML pages"
+    )
+
+
+class URLLoader:
+    """
+    Loads documents from URLs using scraper and PDF parser plugins.
+
+    Orchestrates the content acquisition process:
+    1. Fetch content using a scraper plugin
+    2. Detect content type (HTML vs PDF)
+    3. Parse content using appropriate parser
+    4. Create a Document object
+
+    Example:
+        >>> loader = URLLoader()
+        >>> document = await loader.load("https://example.com/article")
+        >>> print(document.title)
+
+        >>> # Synchronous usage
+        >>> document = loader.load_sync("https://example.com/report.pdf")
+    """
+
+    def __init__(self, config: Optional[URLLoaderConfig] = None):
+        """
+        Initialize the URL loader.
+
+        Args:
+            config: Loader configuration
+        """
+        self.config = config or URLLoaderConfig()
+        self._scraper: Optional[BaseScraperPlugin] = None
+        self._pdf_parser: Optional[BasePDFParserPlugin] = None
+
+    def _get_scraper(self) -> BaseScraperPlugin:
+        """Get the scraper plugin to use."""
+        if self._scraper is not None:
+            return self._scraper
+
+        scrapers = PluginRegistry.get_scrapers()
+        if not scrapers:
+            raise RuntimeError(
+                "No scraper plugins registered. "
+                "Ensure plugins are loaded via 'from statement_extractor import plugins'"
+            )
+
+        if self.config.scraper_plugin:
+            for scraper in scrapers:
+                if scraper.name == self.config.scraper_plugin:
+                    self._scraper = scraper
+                    return scraper
+            raise ValueError(f"Scraper plugin not found: {self.config.scraper_plugin}")
+
+        # Use first available (highest priority)
+        self._scraper = scrapers[0]
+        return self._scraper
+
+    def _get_pdf_parser(self) -> BasePDFParserPlugin:
+        """Get the PDF parser plugin to use."""
+        if self._pdf_parser is not None:
+            return self._pdf_parser
+
+        parsers = PluginRegistry.get_pdf_parsers()
+        if not parsers:
+            raise RuntimeError(
+                "No PDF parser plugins registered. "
+                "Ensure plugins are loaded via 'from statement_extractor import plugins'"
+            )
+
+        if self.config.pdf_parser_plugin:
+            for parser in parsers:
+                if parser.name == self.config.pdf_parser_plugin:
+                    self._pdf_parser = parser
+                    return parser
+            raise ValueError(f"PDF parser plugin not found: {self.config.pdf_parser_plugin}")
+
+        # Use first available (highest priority)
+        self._pdf_parser = parsers[0]
+        return self._pdf_parser
+
+    async def load(self, url: str) -> Document:
+        """
+        Load a URL and return a Document.
+
+        Args:
+            url: URL to load
+
+        Returns:
+            Document with extracted content
+
+        Raises:
+            ValueError: If URL cannot be fetched or parsed
+        """
+        logger.info(f"Loading URL: {url}")
+
+        # 1. Fetch content
+        scraper = self._get_scraper()
+        result = await scraper.fetch(url, self.config.timeout)
+
+        if not result.ok:
+            raise ValueError(f"Failed to fetch {url}: {result.error}")
+
+        logger.debug(f"Fetched {len(result.content)} bytes, type: {result.content_type}")
+
+        # 2. Process based on content type
+        if result.content_type == ContentType.PDF:
+            return self._process_pdf(result)
+        elif result.content_type == ContentType.HTML:
+            return self._process_html(result)
+        else:
+            # Try to guess based on content
+            if result.content[:5] == b"%PDF-":
+                return self._process_pdf(result)
+            # Default to HTML
+            return self._process_html(result)
+
+    def load_sync(self, url: str) -> Document:
+        """
+        Synchronous wrapper for load().
+
+        Args:
+            url: URL to load
+
+        Returns:
+            Document with extracted content
+        """
+        return asyncio.run(self.load(url))
+
+    def _process_pdf(self, result: ScraperResult) -> Document:
+        """
+        Convert PDF to Document with pages.
+
+        Args:
+            result: ScraperResult containing PDF bytes
+
+        Returns:
+            Document with PDF content
+        """
+        logger.info(f"Processing PDF from {result.final_url}")
+
+        parser = self._get_pdf_parser()
+        parse_result = parser.parse(
+            result.content,
+            max_pages=self.config.max_pdf_pages,
+            use_ocr=self.config.use_ocr,
+        )
+
+        if not parse_result.ok:
+            raise ValueError(f"Failed to parse PDF: {parse_result.error}")
+
+        logger.info(f"Extracted {len(parse_result.pages)} pages from PDF")
+
+        # Create Document from pages
+        kwargs = {
+            "pages": parse_result.pages,
+            "title": parse_result.metadata.get("title"),
+            "source_type": "pdf",
+            "url": result.final_url,
+        }
+        author = parse_result.metadata.get("author")
+        if author:
+            kwargs["authors"] = [author]
+
+        return Document.from_pages(**kwargs)
+
+    def _process_html(self, result: ScraperResult) -> Document:
+        """
+        Convert HTML to Document (single page).
+
+        Args:
+            result: ScraperResult containing HTML bytes
+
+        Returns:
+            Document with HTML content
+        """
+        logger.info(f"Processing HTML from {result.final_url}")
+
+        # Decode HTML
+        try:
+            html = result.content.decode("utf-8", errors="replace")
+        except Exception as e:
+            raise ValueError(f"Failed to decode HTML: {e}")
+
+        # Extract text and metadata
+        if self.config.extract_metadata:
+            text, metadata = extract_article_content(html)
+            title = metadata.get("title")
+            author = metadata.get("author")
+            # Log extracted metadata
+            logger.debug(f"Extracted metadata: {metadata}")
+        else:
+            text, title = extract_text_from_html(html)
+            author = None
+            metadata = {}
+
+        if not text or len(text.strip()) < 50:
+            raise ValueError("No meaningful content extracted from HTML")
+
+        logger.info(f"Extracted {len(text)} chars from HTML")
+        if title:
+            logger.info(f"  Title: {title}")
+        if author:
+            logger.info(f"  Author: {author}")
+        if metadata.get("published_date"):
+            logger.info(f"  Published: {metadata.get('published_date')}")
+
+        # Create Document using from_pages since from_text forces source_type="text"
+        kwargs = {
+            "pages": [text],
+            "title": title,
+            "source_type": "webpage",
+            "url": result.final_url,
+        }
+        if author:
+            kwargs["authors"] = [author]
+
+        return Document.from_pages(**kwargs)
+
+
+async def load_url(
+    url: str,
+    config: Optional[URLLoaderConfig] = None,
+) -> Document:
+    """
+    Convenience function to load a URL.
+
+    Args:
+        url: URL to load
+        config: Optional loader configuration
+
+    Returns:
+        Document with extracted content
+    """
+    loader = URLLoader(config)
+    return await loader.load(url)
+
+
+def load_url_sync(
+    url: str,
+    config: Optional[URLLoaderConfig] = None,
+) -> Document:
+    """
+    Convenience function to load a URL synchronously.
+
+    Args:
+        url: URL to load
+        config: Optional loader configuration
+
+    Returns:
+        Document with extracted content
+    """
+    loader = URLLoader(config)
+    return loader.load_sync(url)

statement_extractor/document/pipeline.py

@@ -0,0 +1,388 @@
+"""
+DocumentPipeline - Orchestrates document-level extraction with chunking and citations.
+
+Wraps ExtractionPipeline to provide document-level features:
+- Text chunking with page awareness
+- Batch processing through pipeline stages
+- Statement deduplication across chunks
+- Citation generation from document metadata
+"""
+
+import logging
+import time
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from ..models.document import ChunkingConfig, Document
+from ..pipeline import ExtractionPipeline, PipelineConfig
+from .chunker import DocumentChunker
+from .context import DocumentContext
+from .deduplicator import StatementDeduplicator
+from .summarizer import DocumentSummarizer
+
+logger = logging.getLogger(__name__)
+
+
+class DocumentPipelineConfig(BaseModel):
+    """Configuration for document pipeline processing."""
+
+    chunking: ChunkingConfig = Field(
+        default_factory=ChunkingConfig,
+        description="Configuration for text chunking"
+    )
+    generate_summary: bool = Field(
+        default=True,
+        description="Whether to generate a document summary"
+    )
+    deduplicate_across_chunks: bool = Field(
+        default=True,
+        description="Whether to deduplicate statements across chunks"
+    )
+    batch_size: int = Field(
+        default=10,
+        ge=1,
+        description="Number of items to process in each batch"
+    )
+    pipeline_config: Optional[PipelineConfig] = Field(
+        default=None,
+        description="Configuration for the underlying extraction pipeline"
+    )
+
+
+class DocumentPipeline:
+    """
+    Document-level extraction pipeline.
+
+    Processes documents through:
+    1. Summary generation (optional)
+    2. Chunking with page awareness
+    3. Batch extraction through all pipeline stages
+    4. Deduplication across chunks
+    5. Citation generation
+
+    Example:
+        >>> pipeline = DocumentPipeline()
+        >>> document = Document.from_text("Long document text...", title="Report")
+        >>> ctx = pipeline.process(document)
+        >>> for stmt in ctx.labeled_statements:
+        ...     print(f"{stmt}: {stmt.citation}")
+    """
+
+    def __init__(
+        self,
+        config: Optional[DocumentPipelineConfig] = None,
+    ):
+        """
+        Initialize the document pipeline.
+
+        Args:
+            config: Document pipeline configuration
+        """
+        self.config = config or DocumentPipelineConfig()
+
+        # Initialize components
+        self._chunker = DocumentChunker(self.config.chunking)
+        self._deduplicator = StatementDeduplicator()
+        self._summarizer = DocumentSummarizer() if self.config.generate_summary else None
+        self._pipeline = ExtractionPipeline(self.config.pipeline_config)
+
+    def process(self, document: Document) -> DocumentContext:
+        """
+        Process a document through the pipeline.
+
+        Args:
+            document: Document to process
+
+        Returns:
+            DocumentContext with all extraction results
+        """
+        logger.info(f"Starting document pipeline: {document.document_id}")
+        start_time = time.time()
+
+        ctx = DocumentContext(document=document)
+
+        try:
+            # Step 1: Generate summary (if enabled)
+            if self.config.generate_summary and self._summarizer:
+                self._generate_summary(document, ctx)
+
+            # Step 2: Chunk the document
+            chunks = self._chunker.chunk_document(document)
+            ctx.chunks = chunks
+            logger.info(f"Created {len(chunks)} chunks")
+
+            if not chunks:
+                logger.warning("No chunks created from document")
+                return ctx
+
+            # Step 3: Process all chunks through Stage 1 (Splitting)
+            self._process_stage1(ctx)
+
+            # Step 4: Deduplicate raw triples
+            if self.config.deduplicate_across_chunks:
+                self._deduplicate_triples(ctx)
+
+            # Step 5: Process through remaining stages (2-6)
+            self._process_remaining_stages(ctx)
+
+            # Step 6: Add citations to statements
+            self._add_citations(ctx)
+
+        except Exception as e:
+            logger.exception("Document pipeline failed")
+            ctx.add_error(f"Pipeline error: {str(e)}")
+            raise
+
+        total_time = time.time() - start_time
+        ctx.record_timing("total", total_time)
+
+        logger.info(
+            f"Document pipeline complete: {ctx.statement_count} statements, "
+            f"{ctx.duplicates_removed} duplicates removed, "
+            f"{total_time:.2f}s"
+        )
+
+        return ctx
+
+    def _generate_summary(self, document: Document, ctx: DocumentContext) -> None:
+        """Generate document summary."""
+        logger.info("Generating document summary")
+        start_time = time.time()
+
+        try:
+            summary = self._summarizer.summarize(document)
+            document.summary = summary
+            logger.info(f"Generated summary: {len(summary)} chars")
+        except Exception as e:
+            logger.error(f"Summary generation failed: {e}")
+            ctx.add_warning(f"Summary generation failed: {e}")
+
+        ctx.record_timing("summarization", time.time() - start_time)
+
+    def _process_stage1(self, ctx: DocumentContext) -> None:
+        """Process all chunks through Stage 1 (Splitting) using batch processing."""
+        from ..pipeline.registry import PluginRegistry
+        from ..plugins.base import PluginCapability
+
+        logger.info(f"Processing {len(ctx.chunks)} chunks through Stage 1 (batch mode)")
+        start_time = time.time()
+
+        # Get the splitter plugin
+        splitters = PluginRegistry.get_splitters()
+        if not splitters:
+            logger.warning("No splitter plugins registered")
+            return
+
+        # Use first enabled splitter
+        splitter = None
+        for s in splitters:
+            plugin_enabled = (
+                self.config.pipeline_config is None or
+                self.config.pipeline_config.is_plugin_enabled(s.name)
+            )
+            if plugin_enabled:
+                splitter = s
+                break
+
+        if not splitter:
+            logger.warning("No enabled splitter plugin found")
+            return
+
+        # Extract all chunk texts
+        chunk_texts = [chunk.text for chunk in ctx.chunks]
+
+        # Create a dummy context for the splitter
+        from ..pipeline.context import PipelineContext
+        dummy_ctx = PipelineContext(
+            source_text="",  # Not used for batch splitting
+            source_metadata=self.config.pipeline_config.model_dump() if self.config.pipeline_config else {},
+        )
+
+        all_triples = []
+
+        # Require batch processing capability
+        if PluginCapability.BATCH_PROCESSING not in splitter.capabilities:
+            raise RuntimeError(
+                f"Splitter plugin '{splitter.name}' does not support batch processing. "
+                "Document pipeline requires BATCH_PROCESSING capability for efficient GPU utilization."
+            )
+
+        logger.info(f"Using batch splitting with {splitter.name}")
+        batch_results = splitter.split_batch(chunk_texts, dummy_ctx)
+
+        # Annotate triples with document/chunk info
+        for chunk, triples in zip(ctx.chunks, batch_results):
+            for triple in triples:
+                triple.document_id = ctx.document.document_id
+                triple.chunk_index = chunk.chunk_index
+                triple.page_number = chunk.primary_page
+                all_triples.append(triple)
+
+        ctx.raw_triples = all_triples
+        ctx.pre_dedup_count = len(all_triples)
+
+        ctx.record_timing("stage1_batch", time.time() - start_time)
+        logger.info(f"Stage 1 produced {len(all_triples)} raw triples from {len(ctx.chunks)} chunks")
+
+    def _deduplicate_triples(self, ctx: DocumentContext) -> None:
+        """Deduplicate raw triples across chunks."""
+        logger.info(f"Deduplicating {len(ctx.raw_triples)} triples")
+        start_time = time.time()
+
+        original_count = len(ctx.raw_triples)
+        ctx.raw_triples = self._deduplicator.deduplicate_batch(ctx.raw_triples)
+        ctx.post_dedup_count = len(ctx.raw_triples)
+
+        removed = original_count - len(ctx.raw_triples)
+        ctx.record_timing("deduplication", time.time() - start_time)
+        logger.info(f"Removed {removed} duplicate triples")
+
+    def _process_remaining_stages(self, ctx: DocumentContext) -> None:
+        """Process through stages 2-6."""
+        logger.info(f"Processing {len(ctx.raw_triples)} triples through stages 2-6")
+        start_time = time.time()
+
+        # Create a pipeline config for stages 2-6
+        # Exclude enabled_stages from base config to avoid duplicate keyword argument
+        base_config = {}
+        if self.config.pipeline_config:
+            base_config = self.config.pipeline_config.model_dump(exclude={"enabled_stages"})
+        stages_config = PipelineConfig(
+            enabled_stages={2, 3, 4, 5, 6},
+            **base_config
+        )
+
+        # Create a combined context with all raw triples
+        from ..pipeline.context import PipelineContext
+
+        combined_ctx = PipelineContext(
+            source_text=ctx.document.full_text,
+            source_metadata={
+                "document_id": ctx.document.document_id,
+                "title": ctx.document.metadata.title,
+            },
+        )
+        combined_ctx.raw_triples = ctx.raw_triples
+
+        # Run stages 2-6
+        pipeline = ExtractionPipeline(stages_config)
+
+        # We need to manually run stages since we're providing pre-existing triples
+        # Stage 2: Extraction
+        if stages_config.is_stage_enabled(2):
+            combined_ctx = pipeline._run_extraction(combined_ctx)
+
+        # Propagate document info to statements
+        for stmt in combined_ctx.statements:
+            # Find the source triple to get document info
+            for triple in ctx.raw_triples:
+                if triple.source_sentence in stmt.source_text:
+                    stmt.document_id = triple.document_id
+                    stmt.chunk_index = triple.chunk_index
+                    stmt.page_number = triple.page_number
+                    break
+
+        # Stage 3: Qualification
+        if stages_config.is_stage_enabled(3):
+            combined_ctx = pipeline._run_qualification(combined_ctx)
+
+        # Stage 4: Canonicalization
+        if stages_config.is_stage_enabled(4):
+            combined_ctx = pipeline._run_canonicalization(combined_ctx)
+
+        # Stage 5: Labeling
+        if stages_config.is_stage_enabled(5):
+            combined_ctx = pipeline._run_labeling(combined_ctx)
+
+        # Stage 6: Taxonomy
+        if stages_config.is_stage_enabled(6):
+            combined_ctx = pipeline._run_taxonomy(combined_ctx)
+
+        # Propagate document info to labeled statements
+        for labeled_stmt in combined_ctx.labeled_statements:
+            labeled_stmt.document_id = labeled_stmt.statement.document_id
+            labeled_stmt.page_number = labeled_stmt.statement.page_number
+
+        ctx.statements = combined_ctx.statements
+        ctx.labeled_statements = combined_ctx.labeled_statements
+
+        # Merge timings
+        for stage, duration in combined_ctx.stage_timings.items():
+            ctx.record_timing(stage, duration)
+
+        ctx.record_timing("stages_2_6_batch", time.time() - start_time)
+        logger.info(f"Stages 2-6 produced {len(ctx.labeled_statements)} labeled statements")
+
+    def _add_citations(self, ctx: DocumentContext) -> None:
+        """Add citations to all labeled statements."""
+        logger.info("Adding citations to statements")
+
+        for stmt in ctx.labeled_statements:
+            citation = ctx.document.metadata.format_citation(stmt.page_number)
+            stmt.citation = citation if citation else None
+
+    def process_text(
+        self,
+        text: str,
+        title: Optional[str] = None,
+        **metadata_kwargs,
+    ) -> DocumentContext:
+        """
+        Process plain text through the document pipeline.
+
+        Convenience method that creates a Document from text.
+
+        Args:
+            text: Text to process
+            title: Optional document title
+            **metadata_kwargs: Additional document metadata
+
+        Returns:
+            DocumentContext with extraction results
+        """
+        document = Document.from_text(text, title=title, **metadata_kwargs)
+        return self.process(document)
+
+    async def process_url(
+        self,
+        url: str,
+        loader_config: Optional["URLLoaderConfig"] = None,
+    ) -> DocumentContext:
+        """
+        Process a URL through the document pipeline.
+
+        Fetches the URL, extracts content (HTML or PDF), and processes it.
+
+        Args:
+            url: URL to process
+            loader_config: Optional loader configuration
+
+        Returns:
+            DocumentContext with extraction results
+        """
+        from .loader import URLLoader, URLLoaderConfig
+
+        loader = URLLoader(loader_config or URLLoaderConfig())
+        document = await loader.load(url)
+        return self.process(document)
+
+    def process_url_sync(
+        self,
+        url: str,
+        loader_config: Optional["URLLoaderConfig"] = None,
+    ) -> DocumentContext:
+        """
+        Process a URL through the document pipeline (synchronous).
+
+        Fetches the URL, extracts content (HTML or PDF), and processes it.
+
+        Args:
+            url: URL to process
+            loader_config: Optional loader configuration
+
+        Returns:
+            DocumentContext with extraction results
+        """
+        import asyncio
+        return asyncio.run(self.process_url(url, loader_config))