natural-pdf 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

natural_pdf/extraction/manager.py

@@ -0,0 +1,134 @@
+import logging
+from typing import Any, Type, Optional
+from pydantic import BaseModel
+import io
+import base64
+from PIL import Image
+
+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,246 @@
+import logging
+from typing import TYPE_CHECKING, Any, Type, Optional
+from abc import ABC, abstractmethod
+from pydantic import BaseModel
+
+# Avoid circular import
+if TYPE_CHECKING:
+    from natural_pdf.extraction.result import StructuredDataResult
+    from natural_pdf.core.page import Page
+    from natural_pdf.elements.base import Element
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_STRUCTURED_KEY = "default-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,37 @@
+from typing import Optional, TypeVar, Generic, Any
+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,7 +11,13 @@ logger = logging.getLogger("natural_pdf.ocr")
 
 # Import the base classes that are always available
 from .engine import OCREngine
-from .ocr_options import OCROptions, BaseOCROptions, EasyOCROptions, PaddleOCROptions, SuryaOCROptions
+from .ocr_options import (
+    OCROptions,
+    BaseOCROptions,
+    EasyOCROptions,
+    PaddleOCROptions,
+    SuryaOCROptions,
+)
 from .ocr_manager import OCRManager
 from .ocr_factory import OCRFactory
 
@@ -22,13 +28,14 @@ __all__ = [
     "OCROptions",
     "BaseOCROptions",
     "EasyOCROptions",
-    "PaddleOCROptions", 
+    "PaddleOCROptions",
     "SuryaOCROptions",
     "OCRFactory",
     "get_engine",
-    "list_available_engines"
+    "list_available_engines",
 ]
 
+
 def get_engine(engine_name=None, **kwargs):
     """
     Get OCR engine by name with graceful handling of missing dependencies.
@@ -40,27 +47,27 @@ def get_engine(engine_name=None, **kwargs):
 
     Returns:
         OCREngine instance
-        
+
     Raises:
         ImportError: If the requested engine's dependencies aren't installed
         ValueError: If the engine_name is unknown
     """
     logger.debug(f"Initializing OCR engine: {engine_name or 'best available'}")
-    
+
     try:
         if engine_name is None or engine_name == "default":
             # Use the factory to get the best available engine
             engine = OCRFactory.get_recommended_engine(**kwargs)
             logger.info(f"Using recommended OCR engine: {engine.__class__.__name__}")
             return engine
-            
+
         # Use the factory to create a specific engine
         normalized_name = engine_name.lower()
         if normalized_name in ["easyocr", "paddle", "surya"]:
             return OCRFactory.create_engine(normalized_name, **kwargs)
         else:
             raise ValueError(f"Unknown OCR engine: {engine_name}")
-            
+
     except ImportError as e:
         logger.error(f"OCR engine dependency error: {e}")
         raise
@@ -68,10 +75,11 @@ def get_engine(engine_name=None, **kwargs):
         logger.error(f"Error initializing OCR engine: {e}")
         raise
 
+
 def list_available_engines():
     """
     List all available OCR engines.
-    
+
     Returns:
         Dict[str, bool]: Dictionary mapping engine names to availability status
     """

natural_pdf/ocr/engine.py

@@ -13,11 +13,17 @@ logger = logging.getLogger(__name__)
 
 class TextRegion:
     """Standard representation of an OCR text region."""
-    
-    def __init__(self, bbox: Tuple[float, float, float, float], text: str, confidence: float, source: str = "ocr"):
+
+    def __init__(
+        self,
+        bbox: Tuple[float, float, float, float],
+        text: str,
+        confidence: float,
+        source: str = "ocr",
+    ):
         """
         Initialize a text region.
-        
+
         Args:
             bbox: Tuple of (x0, y0, x1, y1) coordinates
             text: The recognized text
@@ -28,7 +34,7 @@ class TextRegion:
         self.text = text
         self.confidence = confidence
         self.source = source
-    
+
     @classmethod
     def from_polygon(cls, polygon: List[List[float]], text: str, confidence: float):
         """Create from polygon coordinates [[x1,y1], [x2,y2], ...]"""
@@ -36,24 +42,24 @@ class TextRegion:
         y_coords = [float(point[1]) for point in polygon]
         bbox = (min(x_coords), min(y_coords), max(x_coords), max(y_coords))
         return cls(bbox, text, confidence)
-    
+
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary representation for compatibility."""
         return {
             "bbox": self.bbox,
             "text": self.text,
             "confidence": self.confidence,
-            "source": self.source
+            "source": self.source,
         }
 
 
 class OCREngine(ABC):
     """Abstract Base Class for OCR engines."""
-    
+
     # Default values as class constants
     DEFAULT_MIN_CONFIDENCE = 0.2
-    DEFAULT_LANGUAGES = ['en']
-    DEFAULT_DEVICE = 'cpu'
+    DEFAULT_LANGUAGES = ["en"]
+    DEFAULT_DEVICE = "cpu"
 
     def __init__(self):
         """Initializes the base OCR engine."""
@@ -74,7 +80,7 @@ class OCREngine(ABC):
     ) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:
         """
         Process a single image or batch of images with OCR.
-        
+
         Args:
             images: A single PIL Image or a list of PIL Images
             languages: List of languages to use (default: ['en'])
@@ -82,7 +88,7 @@ class OCREngine(ABC):
             device: Device to use for processing (default: 'cpu')
             detect_only: Whether to only detect text regions without recognition
             options: Engine-specific options
-            
+
         Returns:
             For a single image: List of text region dictionaries
             For a batch: List of lists of text region dictionaries
@@ -90,42 +96,48 @@ class OCREngine(ABC):
         # Convert single image to batch format
         single_image = not isinstance(images, list)
         image_batch = [images] if single_image else images
-        
+
         # Use default values where parameters are not provided
         effective_languages = languages or self.DEFAULT_LANGUAGES
-        effective_confidence = min_confidence if min_confidence is not None else self.DEFAULT_MIN_CONFIDENCE
+        effective_confidence = (
+            min_confidence if min_confidence is not None else self.DEFAULT_MIN_CONFIDENCE
+        )
         effective_device = device or self.DEFAULT_DEVICE
-        
+
         # Ensure the model is initialized
         self._ensure_initialized(effective_languages, effective_device, options)
-        
+
         # Process each image in the batch
         results = []
         for img in image_batch:
             # Preprocess the image for the specific engine
             processed_img = self._preprocess_image(img)
-            
+
             # Process the image with the engine-specific implementation
             raw_results = self._process_single_image(processed_img, detect_only, options)
-            
+
             # Convert results to standardized format
             text_regions = self._standardize_results(raw_results, effective_confidence, detect_only)
-            
+
             # Convert TextRegion objects to dictionaries for backward compatibility
             region_dicts = [region.to_dict() for region in text_regions]
             results.append(region_dicts)
-        
+
         # Return results in the appropriate format
         return results[0] if single_image else results
 
-    def _ensure_initialized(self, languages: List[str], device: str, options: Optional[BaseOCROptions]):
+    def _ensure_initialized(
+        self, languages: List[str], device: str, options: Optional[BaseOCROptions]
+    ):
         """Ensure the model is initialized with the correct parameters."""
         if not self._initialized:
             self._initialize_model(languages, device, options)
             self._initialized = True
-            
+
     @abstractmethod
-    def _initialize_model(self, languages: List[str], device: str, options: Optional[BaseOCROptions]):
+    def _initialize_model(
+        self, languages: List[str], device: str, options: Optional[BaseOCROptions]
+    ):
         """Initialize the OCR model with the given parameters."""
         raise NotImplementedError("Subclasses must implement this method")
 
@@ -133,14 +145,18 @@ class OCREngine(ABC):
     def _preprocess_image(self, image: Image.Image) -> Any:
         """Convert PIL Image to engine-specific format."""
         raise NotImplementedError("Subclasses must implement this method")
-        
+
     @abstractmethod
-    def _process_single_image(self, image: Any, detect_only: bool, options: Optional[BaseOCROptions]) -> Any:
+    def _process_single_image(
+        self, image: Any, detect_only: bool, options: Optional[BaseOCROptions]
+    ) -> Any:
         """Process a single image with the initialized model."""
         raise NotImplementedError("Subclasses must implement this method")
-        
+
     @abstractmethod
-    def _standardize_results(self, raw_results: Any, min_confidence: float, detect_only: bool) -> List[TextRegion]:
+    def _standardize_results(
+        self, raw_results: Any, min_confidence: float, detect_only: bool
+    ) -> List[TextRegion]:
         """Convert engine-specific results to standardized TextRegion objects."""
         raise NotImplementedError("Subclasses must implement this method")
 
@@ -181,23 +197,23 @@ class OCREngine(ABC):
                 return tuple(float(c) for c in bbox[:4])
             except (ValueError, TypeError) as e:
                 raise ValueError(f"Invalid number format in bbox: {bbox}") from e
-        
+
         # Check if it's in polygon format [[x1,y1],[x2,y2],...]
         elif (
             isinstance(bbox, (list, tuple))
             and len(bbox) > 0
             and isinstance(bbox[0], (list, tuple))
-            and len(bbox[0]) == 2 # Ensure points are pairs
+            and len(bbox[0]) == 2  # Ensure points are pairs
         ):
             try:
                 x_coords = [float(point[0]) for point in bbox]
                 y_coords = [float(point[1]) for point in bbox]
-                if not x_coords or not y_coords: # Handle empty polygon case
+                if not x_coords or not y_coords:  # Handle empty polygon case
                     raise ValueError("Empty polygon provided")
                 return (min(x_coords), min(y_coords), max(x_coords), max(y_coords))
             except (ValueError, TypeError, IndexError) as e:
                 raise ValueError(f"Invalid polygon format or values: {bbox}") from e
-        
+
         # If it's neither format, raise an error
         raise ValueError(f"Could not standardize bounding box from unexpected format: {bbox}")