natural-pdf 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

natural_pdf/elements/text.py

@@ -276,7 +276,10 @@ class TextElement(Element):
 
     def __repr__(self) -> str:
         """String representation of the text element."""
-        preview = self.text[:10] + "..." if len(self.text) > 10 else self.text
+        if self.text:
+            preview = self.text[:10] + "..." if len(self.text) > 10 else self.text
+        else:
+            preview = "..."
         font_style = []
         if self.bold:
             font_style.append("bold")

natural_pdf/export/mixin.py

@@ -0,0 +1,137 @@
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+logger = logging.getLogger(__name__)
+
+
+class ExportMixin:
+    """
+    Mixin for exporting analyses from collections of elements.
+
+    This mixin is designed to be used with PDF, PDFCollection,
+    PageCollection, and ElementCollection classes.
+    """
+
+    def export_analyses(
+        self,
+        output_path: str,
+        analysis_keys: Union[str, List[str]],
+        format: str = "json",
+        include_content: bool = True,
+        include_images: bool = False,
+        image_dir: Optional[str] = None,
+        image_format: str = "jpg",
+        image_resolution: int = 72,
+        overwrite: bool = True,
+        **kwargs,
+    ) -> str:
+        """
+        Export analysis results to a file.
+
+        Args:
+            output_path: Path to save the export file
+            analysis_keys: Key(s) in the analyses dictionary to export
+            format: Export format ('json', 'csv', 'excel')
+            include_content: Whether to include extracted text
+            include_images: Whether to export images of elements
+            image_dir: Directory to save images (created if doesn't exist)
+            image_format: Format to save images ('jpg', 'png')
+            image_resolution: Resolution for exported images
+            overwrite: Whether to overwrite existing files
+            **kwargs: Additional format-specific options
+
+        Returns:
+            Path to the exported file
+        """
+        # Convert single key to list for consistency
+        if isinstance(analysis_keys, str):
+            analysis_keys = [analysis_keys]
+
+        # Create output directory
+        output_path = Path(output_path)
+        os.makedirs(output_path.parent, exist_ok=True)
+
+        # Check if file exists and handle overwrite
+        if output_path.exists() and not overwrite:
+            raise FileExistsError(f"Output file {output_path} already exists and overwrite=False")
+
+        # Prepare image directory if needed
+        if include_images:
+            if image_dir is None:
+                image_dir = output_path.parent / f"{output_path.stem}_images"
+            os.makedirs(image_dir, exist_ok=True)
+            image_dir = Path(image_dir)  # Convert to Path object
+
+        # Gather data from collection
+        data = self._gather_analysis_data(
+            analysis_keys=analysis_keys,
+            include_content=include_content,
+            include_images=include_images,
+            image_dir=image_dir,
+            image_format=image_format,
+            image_resolution=image_resolution,
+        )
+
+        # Export based on format
+        if format.lower() == "json":
+            return self._export_to_json(data, output_path, **kwargs)
+        elif format.lower() == "csv":
+            return self._export_to_csv(data, output_path, **kwargs)
+        elif format.lower() == "excel":
+            return self._export_to_excel(data, output_path, **kwargs)
+        else:
+            raise ValueError(f"Unsupported export format: {format}")
+
+    def _gather_analysis_data(
+        self,
+        analysis_keys: List[str],
+        include_content: bool,
+        include_images: bool,
+        image_dir: Optional[Path],
+        image_format: str,
+        image_resolution: int,
+    ) -> List[Dict[str, Any]]:
+        """
+        Gather analysis data from elements in the collection.
+
+        This method should be implemented by each collection class.
+        """
+        raise NotImplementedError("Subclasses must implement _gather_analysis_data")
+
+    def _export_to_json(self, data: List[Dict[str, Any]], output_path: Path, **kwargs) -> str:
+        """Export data to JSON format."""
+        with open(output_path, "w") as f:
+            json.dump(data, f, indent=2, **kwargs)
+        logger.info(f"Exported analysis data to {output_path}")
+        return str(output_path)
+
+    def _export_to_csv(self, data: List[Dict[str, Any]], output_path: Path, **kwargs) -> str:
+        """Export data to CSV format."""
+        try:
+            import pandas as pd
+
+            # Normalize nested data
+            df = pd.json_normalize(data)
+            df.to_csv(output_path, index=False, **kwargs)
+            logger.info(f"Exported analysis data to {output_path}")
+            return str(output_path)
+        except ImportError:
+            raise ImportError("Pandas is required for CSV export. Install with: pip install pandas")
+
+    def _export_to_excel(self, data: List[Dict[str, Any]], output_path: Path, **kwargs) -> str:
+        """Export data to Excel format."""
+        try:
+            import pandas as pd
+
+            # Normalize nested data
+            df = pd.json_normalize(data)
+            df.to_excel(output_path, index=False, **kwargs)
+            logger.info(f"Exported analysis data to {output_path}")
+            return str(output_path)
+        except ImportError:
+            raise ImportError(
+                "Pandas and openpyxl are required for Excel export. Install with: pip install pandas openpyxl"
+            )

natural_pdf/exporters/base.py

@@ -1,10 +1,10 @@
 import abc
 import logging
-from typing import Union, List, TYPE_CHECKING
+from typing import TYPE_CHECKING, List, Union
 
 if TYPE_CHECKING:
-    from natural_pdf.core.pdf import PDF
     from natural_pdf.collections.pdf_collection import PDFCollection
+    from natural_pdf.core.pdf import PDF
 
 logger = logging.getLogger(__name__)
 
@@ -40,8 +40,8 @@ class FinetuneExporter(abc.ABC):
         """
         Helper to consistently resolve the input source to a list of PDF objects.
         """
-        from natural_pdf.core.pdf import PDF  # Avoid circular import at module level
         from natural_pdf.collections.pdf_collection import PDFCollection  # Avoid circular import
+        from natural_pdf.core.pdf import PDF  # Avoid circular import at module level
 
         pdfs_to_process: List["PDF"] = []
         if isinstance(source, PDF):

natural_pdf/exporters/paddleocr.py

@@ -1,8 +1,9 @@
-import os
 import logging
+import os
 import random
 import shutil
-from typing import Union, List, Optional, TYPE_CHECKING, Set, Tuple
+from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union
+
 from tqdm import tqdm
 
 from natural_pdf.exporters.base import FinetuneExporter
@@ -11,8 +12,8 @@ from natural_pdf.exporters.base import FinetuneExporter
 from natural_pdf.utils.identifiers import generate_short_path_hash
 
 if TYPE_CHECKING:
-    from natural_pdf.core.pdf import PDF
     from natural_pdf.collections.pdf_collection import PDFCollection
+    from natural_pdf.core.pdf import PDF
     from natural_pdf.elements.text import TextElement
 
 logger = logging.getLogger(__name__)
@@ -48,7 +49,7 @@ class PaddleOCRRecognitionExporter(FinetuneExporter):
             selector: CSS-like selector to filter which TextElements to export.
                       If None and corrected_only is False, all 'text' elements are considered.
             corrected_only: If True, overrides selector and exports only elements likely
-                            originating from a correction manifest (selector="text[source^=manifest]").
+                            originating from a correction manifest (selector="text[source=manifest]").
                             (default: False).
             split_ratio: Ratio for splitting data into training/validation sets (e.g., 0.9 for 90% train).
                          If None, creates a single `label.txt` file (default: 0.9).

natural_pdf/extraction/manager.py

@@ -0,0 +1,135 @@
+import base64
+import io
+import logging
+from typing import Any, Optional, Type
+
+from PIL import Image
+from pydantic import BaseModel
+
+from natural_pdf.extraction.result import StructuredDataResult
+
+logger = logging.getLogger(__name__)
+
+
+class StructuredDataManager:
+    """
+    Manages the process of extracting structured data from elements using LLMs.
+
+    This manager is typically accessed via `pdf.get_manager('structured_data')`.
+    It is stateless and relies on parameters passed during method calls.
+    """
+
+    DEFAULT_TEXT_MODEL = "gpt-4o-mini"
+    DEFAULT_VISION_MODEL = "gpt-4o"
+
+    def __init__(self):
+        """Initializes the manager."""
+        logger.info("Initialized StructuredDataManager.")
+
+    def is_available(self) -> bool:
+        """Checks if necessary dependencies are available."""
+        try:
+            import pydantic
+
+            return True
+        except ImportError:
+            logger.warning("Pydantic is required for structured data extraction.")
+            return False
+
+    def _prepare_llm_messages(
+        self, content: Any, prompt: Optional[str], using: str, schema: Type[BaseModel]
+    ) -> list:
+        """Prepares the message list for the LLM API call."""
+        system_prompt = (
+            prompt
+            or f"Extract the information corresponding to the fields in the {schema.__name__} schema. Respond only with the structured data."
+        )
+
+        messages = [{"role": "system", "content": system_prompt}]
+
+        if using == "text":
+            messages.append({"role": "user", "content": str(content)})
+        elif using == "vision":
+            if isinstance(content, Image.Image):
+                buffered = io.BytesIO()
+                content.save(buffered, format="PNG")
+                base64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": "Extract information from this image based on the schema.",
+                            },
+                            {
+                                "type": "image_url",
+                                "image_url": {"url": f"data:image/png;base64,{base64_image}"},
+                            },
+                        ],
+                    }
+                )
+            else:
+                raise TypeError(
+                    f"Content must be a PIL Image for using='vision', got {type(content)}"
+                )
+        else:
+            raise ValueError(f"Unsupported value for 'using': {using}")
+
+        return messages
+
+    def extract(
+        self,
+        content: Any,
+        schema: Type[BaseModel],
+        client: Any,
+        prompt: Optional[str] = None,
+        using: str = "text",
+        model: Optional[str] = None,
+        **kwargs,
+    ) -> StructuredDataResult:
+        """
+        Extract structured data from content using an LLM.
+
+        Args:
+            content: Text string or Image object
+            schema: Pydantic model class for the desired structure
+            client: Initialized LLM client (e.g., OpenAI client)
+            prompt: Optional user-provided instructions
+            using: Modality ('text' or 'vision')
+            model: Specific LLM model identifier
+            **kwargs: Additional parameters for the LLM API call
+
+        Returns:
+            StructuredDataResult object
+        """
+        logger.debug(f"Extract request: using='{using}', schema='{schema.__name__}'")
+
+        if isinstance(content, list) and using == "vision":
+            if len(content) == 1:
+                content = content[0]
+            elif len(content) > 1:
+                logger.error("Vision extraction not supported for multi-page PDFs")
+                raise NotImplementedError(
+                    "Batch image extraction on multi-page PDF objects is not supported. Apply to individual pages or regions instead."
+                )
+
+        selected_model = model or (
+            self.DEFAULT_VISION_MODEL if using == "vision" else self.DEFAULT_TEXT_MODEL
+        )
+        messages = self._prepare_llm_messages(content, prompt, using, schema)
+
+        try:
+            logger.debug(f"Extracting with model '{selected_model}'")
+            completion = client.beta.chat.completions.parse(
+                model=selected_model, messages=messages, response_format=schema, **kwargs
+            )
+            parsed_data = completion.choices[0].message.parsed
+            return StructuredDataResult(
+                data=parsed_data, success=True, error_message=None, model=selected_model
+            )
+        except Exception as e:
+            logger.error(f"Extraction failed: {str(e)}")
+            return StructuredDataResult(
+                data=None, success=False, error_message=str(e), model=selected_model
+            )

natural_pdf/extraction/mixin.py

@@ -0,0 +1,279 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Optional, Type
+
+from pydantic import BaseModel
+
+# Avoid circular import
+if TYPE_CHECKING:
+    from natural_pdf.core.page import Page
+    from natural_pdf.elements.base import Element
+    from natural_pdf.extraction.result import StructuredDataResult
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_STRUCTURED_KEY = "structured"  # Define default key
+
+
+class ExtractionMixin(ABC):
+    """
+    Mixin class providing structured data extraction capabilities to elements.
+    Assumes the inheriting class has `extract_text(**kwargs)` and `to_image(**kwargs)` methods.
+    """
+
+    def _get_extraction_content(self, using: str = "text", **kwargs) -> Any:
+        """
+        Retrieves the content (text or image) for extraction.
+
+        Args:
+            using: 'text' or 'vision'
+            **kwargs: Additional arguments passed to extract_text or to_image
+
+        Returns:
+            str: Extracted text if using='text'
+            PIL.Image.Image: Rendered image if using='vision'
+            None: If content cannot be retrieved
+        """
+        if not hasattr(self, "extract_text") or not callable(self.extract_text):
+            logger.error(f"ExtractionMixin requires 'extract_text' method on {self!r}")
+            return None
+        if not hasattr(self, "to_image") or not callable(self.to_image):
+            logger.error(f"ExtractionMixin requires 'to_image' method on {self!r}")
+            return None
+
+        try:
+            if using == "text":
+                layout = kwargs.pop("layout", True)
+                return self.extract_text(layout=layout, **kwargs)
+            elif using == "vision":
+                resolution = kwargs.pop("resolution", 72)
+                include_highlights = kwargs.pop("include_highlights", False)
+                labels = kwargs.pop("labels", False)
+                return self.to_image(
+                    resolution=resolution,
+                    include_highlights=include_highlights,
+                    labels=labels,
+                    **kwargs,
+                )
+            else:
+                logger.error(f"Unsupported value for 'using': {using}")
+                return None
+        except Exception as e:
+            logger.error(f"Error getting {using} content from {self!r}: {e}")
+            return None
+
+    def extract(
+        self: Any,
+        schema: Type[BaseModel],
+        client: Any,
+        analysis_key: str = DEFAULT_STRUCTURED_KEY,  # Default key
+        prompt: Optional[str] = None,
+        using: str = "text",
+        model: Optional[str] = None,
+        overwrite: bool = False,  # Add overwrite parameter
+        **kwargs,
+    ) -> Any:
+        """
+        Extracts structured data according to the provided schema.
+
+        Results are stored in the element's `analyses` dictionary.
+
+        Args:
+            schema: Pydantic model class defining the desired structure
+            client: Initialized LLM client
+            analysis_key: Key to store the result under in `analyses`. Defaults to "default-structured".
+            prompt: Optional user-provided prompt for the LLM
+            using: Modality ('text' or 'vision')
+            model: Optional specific LLM model identifier
+            overwrite: If True, allow overwriting an existing result at `analysis_key`.
+            **kwargs: Additional parameters for extraction
+
+        Returns:
+            Self for method chaining
+        """
+        if not analysis_key:
+            raise ValueError("analysis_key cannot be empty for extract operation")
+
+        # --- Overwrite Check --- #
+        if not hasattr(self, "analyses") or self.analyses is None:
+            self.analyses = {}
+
+        if analysis_key in self.analyses and not overwrite:
+            raise ValueError(
+                f"Analysis key '{analysis_key}' already exists in analyses. "
+                f"Use overwrite=True to replace it. Available keys: {list(self.analyses.keys())}"
+            )
+        # --- End Overwrite Check --- #
+
+        # Determine PDF instance to get manager
+        pdf_instance = None
+
+        if hasattr(self, "get_manager") and callable(self.get_manager):
+            # Handle case where self is the PDF instance itself
+            pdf_instance = self
+            logger.debug(f"Manager access via self ({type(self).__name__})")
+        elif (
+            hasattr(self, "pdf")
+            and hasattr(self.pdf, "get_manager")
+            and callable(self.pdf.get_manager)
+        ):
+            # Handle Page or other elements with direct .pdf reference
+            pdf_instance = self.pdf
+            logger.debug(f"Manager access via self.pdf ({type(self).__name__})")
+        elif (
+            hasattr(self, "page")
+            and hasattr(self.page, "pdf")
+            and hasattr(self.page.pdf, "get_manager")
+            and callable(self.page.pdf.get_manager)
+        ):
+            # Handle Region or other elements with .page.pdf reference
+            pdf_instance = self.page.pdf
+            logger.debug(f"Manager access via self.page.pdf ({type(self).__name__})")
+        else:
+            logger.error(
+                f"Could not find get_manager on {type(self).__name__}, self.pdf, or self.page.pdf"
+            )
+            raise RuntimeError(
+                f"Cannot access PDF manager: {type(self).__name__} lacks necessary references"
+            )
+
+        try:
+            manager = pdf_instance.get_manager("structured_data")
+        except Exception as e:
+            raise RuntimeError(f"Failed to get StructuredDataManager: {e}")
+
+        if not manager or not manager.is_available():
+            raise RuntimeError("StructuredDataManager is not available")
+
+        # Get content
+        layout_for_text = kwargs.pop("layout", True)
+        content = self._get_extraction_content(
+            using=using, layout=layout_for_text, **kwargs
+        )  # Pass kwargs
+
+        if content is None or (
+            using == "text" and isinstance(content, str) and not content.strip()
+        ):
+            logger.warning(f"No content available for extraction (using='{using}') on {self!r}")
+            # Import here to avoid circularity at module level
+            from natural_pdf.extraction.result import StructuredDataResult
+
+            result = StructuredDataResult(
+                data=None,
+                success=False,
+                error_message=f"No content available for extraction (using='{using}')",
+                model=model,  # Use model requested, even if failed
+            )
+        else:
+            result = manager.extract(
+                content=content,
+                schema=schema,
+                client=client,
+                prompt=prompt,
+                using=using,
+                model=model,
+                **kwargs,
+            )
+
+        # Store the result
+        self.analyses[analysis_key] = result
+        logger.info(
+            f"Stored extraction result under key '{analysis_key}' (Success: {result.success})"
+        )
+
+        return self
+
+    def extracted(
+        self, field_name: Optional[str] = None, analysis_key: Optional[str] = None
+    ) -> Any:
+        """
+        Convenience method to access results from structured data extraction.
+
+        Args:
+            field_name: The specific field to retrieve from the extracted data dictionary.
+                        If None, returns the entire data dictionary.
+            analysis_key: The key under which the extraction result was stored in `analyses`.
+                          If None, defaults to "default-structured".
+
+        Returns:
+            The requested field value, the entire data dictionary, or raises an error.
+
+        Raises:
+            KeyError: If the specified `analysis_key` is not found in `analyses`.
+            ValueError: If the stored result for `analysis_key` indicates a failed extraction.
+            AttributeError: If the element does not have an `analyses` attribute.
+            KeyError: (Standard Python) If `field_name` is specified but not found in the data.
+        """
+        target_key = analysis_key if analysis_key is not None else DEFAULT_STRUCTURED_KEY
+
+        if not hasattr(self, "analyses") or self.analyses is None:
+            raise AttributeError(f"{type(self).__name__} object has no 'analyses' attribute yet.")
+
+        if target_key not in self.analyses:
+            available_keys = list(self.analyses.keys())
+            raise KeyError(
+                f"Extraction '{target_key}' not found in analyses. "
+                f"Available extractions: {available_keys}"
+            )
+
+        # Import here to avoid circularity and allow type checking
+        from natural_pdf.extraction.result import StructuredDataResult
+
+        result: StructuredDataResult = self.analyses[target_key]
+
+        if not isinstance(result, StructuredDataResult):
+            logger.warning(
+                f"Item found at key '{target_key}' is not a StructuredDataResult (type: {type(result)}). Cannot process."
+            )
+            raise TypeError(
+                f"Expected a StructuredDataResult at key '{target_key}', found {type(result).__name__}"
+            )
+
+        if not result.success:
+            raise ValueError(
+                f"Stored result for '{target_key}' indicates a failed extraction attempt. "
+                f"Error: {result.error_message}"
+            )
+
+        if result.data is None:
+            # This case might occur if success=True but data is somehow None
+            raise ValueError(
+                f"Extraction result for '{target_key}' has no data available, despite success flag."
+            )
+
+        if field_name is None:
+            # Return the whole data object (Pydantic model instance or dict)
+            return result.data
+        else:
+            # Try dictionary key access first, then attribute access
+            if isinstance(result.data, dict):
+                try:
+                    return result.data[field_name]
+                except KeyError:
+                    available_keys = list(result.data.keys())
+                    raise KeyError(
+                        f"Field/Key '{field_name}' not found in extracted dictionary "
+                        f"for key '{target_key}'. Available keys: {available_keys}"
+                    )
+            else:
+                # Assume it's an object, try attribute access
+                try:
+                    return getattr(result.data, field_name)
+                except AttributeError:
+                    # Try to get available fields from the object
+                    available_fields = []
+                    if hasattr(result.data, "model_fields"):  # Pydantic v2
+                        available_fields = list(result.data.model_fields.keys())
+                    elif hasattr(result.data, "__fields__"):  # Pydantic v1
+                        available_fields = list(result.data.__fields__.keys())
+                    elif hasattr(result.data, "__dict__"):  # Fallback
+                        available_fields = list(result.data.__dict__.keys())
+
+                    raise AttributeError(
+                        f"Field/Attribute '{field_name}' not found on extracted object of type {type(result.data).__name__} "
+                        f"for key '{target_key}'. Available fields/attributes: {available_fields}"
+                    )
+                except Exception as e:  # Catch other potential errors during getattr
+                    raise TypeError(
+                        f"Could not access field/attribute '{field_name}' on extracted data for key '{target_key}' (type: {type(result.data).__name__}). Error: {e}"
+                    ) from e

natural_pdf/extraction/result.py

@@ -0,0 +1,23 @@
+from typing import Any, Generic, Optional, TypeVar
+
+from pydantic import BaseModel, Field
+
+# Generic type for the Pydantic model used in the schema
+T_Schema = TypeVar("T_Schema", bound=BaseModel)
+
+
+class StructuredDataResult(BaseModel, Generic[T_Schema]):
+    """
+    Represents the result of a structured data extraction operation.
+
+    Contains the extracted data, success status, and error information.
+    """
+
+    data: Optional[T_Schema] = Field(None, description="Validated data model or None on failure")
+    success: bool = Field(..., description="Whether extraction succeeded")
+    error_message: Optional[str] = Field(None, description="Error details if extraction failed")
+    raw_output: Optional[Any] = Field(None, description="Raw output from the language model")
+    model_used: Optional[str] = Field(None, description="Identifier of the language model used")
+
+    class Config:
+        arbitrary_types_allowed = True

natural_pdf/ocr/__init__.py

@@ -11,15 +11,15 @@ logger = logging.getLogger("natural_pdf.ocr")
 
 # Import the base classes that are always available
 from .engine import OCREngine
+from .ocr_factory import OCRFactory
+from .ocr_manager import OCRManager
 from .ocr_options import (
-    OCROptions,
     BaseOCROptions,
     EasyOCROptions,
+    OCROptions,
     PaddleOCROptions,
     SuryaOCROptions,
 )
-from .ocr_manager import OCRManager
-from .ocr_factory import OCRFactory
 
 # Add all public symbols that should be available when importing this module
 __all__ = [
@@ -41,7 +41,7 @@ def get_engine(engine_name=None, **kwargs):
     Get OCR engine by name with graceful handling of missing dependencies.
 
     Args:
-        engine_name: Name of the engine to use ('easyocr', 'paddle', 'surya')
+        engine_name: Name of the engine to use ('easyocr', 'paddle', 'surya', 'doctr')
                      If None, the best available engine is used
         **kwargs: Additional arguments to pass to the engine constructor
 
@@ -63,7 +63,7 @@ def get_engine(engine_name=None, **kwargs):
 
         # Use the factory to create a specific engine
         normalized_name = engine_name.lower()
-        if normalized_name in ["easyocr", "paddle", "surya"]:
+        if normalized_name in ["easyocr", "paddle", "surya", "doctr"]:
             return OCRFactory.create_engine(normalized_name, **kwargs)
         else:
             raise ValueError(f"Unknown OCR engine: {engine_name}")