visual-rag-toolkit 0.1.1__py3-none-any.whl

visual_rag/indexing/cloudinary_uploader.py

@@ -0,0 +1,274 @@
+"""
+Cloudinary Uploader - Upload images to Cloudinary CDN.
+
+Works INDEPENDENTLY of PDF processing and embedding.
+Use it if you just need to upload images to a CDN.
+
+Features:
+- Retry logic with timeouts
+- Batch uploading
+- Automatic JPEG optimization
+
+Environment Variables:
+- VISUAL_RAG_THREAD_SAFE: Set to "1" to use thread-safe timeouts
+  (required for Streamlit, Flask, or other threaded contexts)
+"""
+
+import io
+import logging
+import os
+import platform
+import signal
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FuturesTimeoutError
+from typing import Optional
+
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+THREAD_SAFE_MODE = os.getenv("VISUAL_RAG_THREAD_SAFE", "").lower() in ("1", "true", "yes")
+
+
+class CloudinaryUploader:
+    """
+    Upload images to Cloudinary CDN.
+
+    Works independently - just needs PIL images.
+
+    Args:
+        cloud_name: Cloudinary cloud name
+        api_key: Cloudinary API key
+        api_secret: Cloudinary API secret
+        folder: Base folder for uploads
+        max_retries: Number of retry attempts
+        timeout_seconds: Timeout per upload
+
+    Example:
+        >>> uploader = CloudinaryUploader(
+        ...     cloud_name="my-cloud",
+        ...     api_key="xxx",
+        ...     api_secret="yyy",
+        ...     folder="my-project",
+        ... )
+        >>>
+        >>> url = uploader.upload(image, "doc_page_1")
+        >>> print(url)  # https://res.cloudinary.com/.../doc_page_1.jpg
+    """
+
+    def __init__(
+        self,
+        cloud_name: Optional[str] = None,
+        api_key: Optional[str] = None,
+        api_secret: Optional[str] = None,
+        folder: str = "visual-rag",
+        max_retries: int = 3,
+        timeout_seconds: int = 30,
+        jpeg_quality: int = 95,
+    ):
+        # Load from environment if not provided
+        self.cloud_name = cloud_name or os.getenv("CLOUDINARY_CLOUD_NAME")
+        self.api_key = api_key or os.getenv("CLOUDINARY_API_KEY")
+        self.api_secret = api_secret or os.getenv("CLOUDINARY_API_SECRET")
+
+        if not all([self.cloud_name, self.api_key, self.api_secret]):
+            raise ValueError(
+                "Cloudinary credentials required. Set CLOUDINARY_CLOUD_NAME, "
+                "CLOUDINARY_API_KEY, CLOUDINARY_API_SECRET environment variables "
+                "or pass them as arguments."
+            )
+
+        self.folder = folder
+        self.max_retries = max_retries
+        self.timeout_seconds = timeout_seconds
+        self.jpeg_quality = jpeg_quality
+
+        # Check dependency
+        try:
+            import cloudinary  # noqa
+        except ImportError:
+            raise ImportError(
+                "Cloudinary not installed. "
+                "Install with: pip install visual-rag-toolkit[cloudinary]"
+            )
+
+        logger.info("☁️ Cloudinary uploader initialized")
+        logger.info(f"   Folder: {folder}")
+
+    def upload(
+        self,
+        image: Image.Image,
+        public_id: str,
+        subfolder: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Upload a single image to Cloudinary.
+
+        Args:
+            image: PIL Image to upload
+            public_id: Public ID (filename without extension)
+            subfolder: Optional subfolder within base folder
+
+        Returns:
+            Secure URL of uploaded image, or None if failed
+        """
+        import cloudinary
+        import cloudinary.uploader
+
+        # Prepare buffer
+        buffer = io.BytesIO()
+        image.save(buffer, format="JPEG", quality=self.jpeg_quality, optimize=True)
+
+        # Configure Cloudinary
+        cloudinary.config(
+            cloud_name=self.cloud_name,
+            api_key=self.api_key,
+            api_secret=self.api_secret,
+        )
+
+        # Build folder path
+        folder_path = self.folder
+        if subfolder:
+            folder_path = f"{self.folder}/{subfolder}"
+
+        def do_upload():
+            buffer.seek(0)
+            result = cloudinary.uploader.upload(
+                buffer,
+                folder=folder_path,
+                overwrite=True,
+                public_id=public_id,
+                resource_type="image",
+                timeout=self.timeout_seconds,
+            )
+            return result["secure_url"]
+
+        # Use thread-safe mode for Streamlit/Flask/threaded contexts
+        # Set VISUAL_RAG_THREAD_SAFE=1 to enable
+        if THREAD_SAFE_MODE or threading.current_thread() is not threading.main_thread():
+            return self._upload_with_thread_timeout(do_upload, public_id)
+        else:
+            return self._upload_with_signal_timeout(do_upload, public_id)
+
+    def _upload_with_thread_timeout(self, do_upload, public_id: str) -> Optional[str]:
+        """Thread-safe upload with ThreadPoolExecutor timeout."""
+        for attempt in range(self.max_retries):
+            try:
+                with ThreadPoolExecutor(max_workers=1) as executor:
+                    future = executor.submit(do_upload)
+                    return future.result(timeout=self.timeout_seconds)
+
+            except FuturesTimeoutError:
+                logger.warning(
+                    f"Upload timeout (attempt {attempt + 1}/{self.max_retries}): {public_id}"
+                )
+                if attempt < self.max_retries - 1:
+                    time.sleep(2**attempt)
+
+            except Exception as e:
+                logger.warning(f"Upload failed (attempt {attempt + 1}/{self.max_retries}): {e}")
+                if attempt < self.max_retries - 1:
+                    time.sleep(2**attempt)
+
+        logger.error(f"❌ Upload failed after {self.max_retries} attempts: {public_id}")
+        return None
+
+    def _upload_with_signal_timeout(self, do_upload, public_id: str) -> Optional[str]:
+        """Signal-based upload timeout (main thread only, Unix/macOS)."""
+        use_timeout = platform.system() != "Windows"
+
+        class SignalTimeoutError(Exception):
+            pass
+
+        def timeout_handler(signum, frame):
+            raise SignalTimeoutError(f"Upload timed out after {self.timeout_seconds}s")
+
+        for attempt in range(self.max_retries):
+            try:
+                if use_timeout:
+                    old_handler = signal.signal(signal.SIGALRM, timeout_handler)
+                    signal.alarm(self.timeout_seconds)
+
+                try:
+                    return do_upload()
+                finally:
+                    if use_timeout:
+                        signal.alarm(0)
+                        signal.signal(signal.SIGALRM, old_handler)
+
+            except SignalTimeoutError:
+                logger.warning(
+                    f"Upload timeout (attempt {attempt + 1}/{self.max_retries}): {public_id}"
+                )
+                if attempt < self.max_retries - 1:
+                    time.sleep(2**attempt)
+
+            except Exception as e:
+                logger.warning(f"Upload failed (attempt {attempt + 1}/{self.max_retries}): {e}")
+                if attempt < self.max_retries - 1:
+                    time.sleep(2**attempt)
+
+        logger.error(f"❌ Upload failed after {self.max_retries} attempts: {public_id}")
+        return None
+
+    def upload_original_and_resized(
+        self,
+        original_image: Image.Image,
+        resized_image: Image.Image,
+        base_public_id: str,
+    ) -> tuple:
+        """
+        Upload both original and resized versions.
+
+        Args:
+            original_image: Original PDF page image
+            resized_image: Resized image for ColPali
+            base_public_id: Base public ID (e.g., "doc_page_1")
+
+        Returns:
+            Tuple of (original_url, resized_url) - either can be None on failure
+        """
+        original_url = self.upload(
+            original_image,
+            base_public_id,
+            subfolder="original",
+        )
+
+        resized_url = self.upload(
+            resized_image,
+            base_public_id,
+            subfolder="resized",
+        )
+
+        return original_url, resized_url
+
+    def upload_original_cropped_and_resized(
+        self,
+        original_image: Image.Image,
+        cropped_image: Optional[Image.Image],
+        resized_image: Image.Image,
+        base_public_id: str,
+    ) -> tuple:
+        original_url = self.upload(
+            original_image,
+            base_public_id,
+            subfolder="original",
+        )
+
+        cropped_url = None
+        if cropped_image is not None:
+            cropped_url = self.upload(
+                cropped_image,
+                base_public_id,
+                subfolder="cropped",
+            )
+
+        resized_url = self.upload(
+            resized_image,
+            base_public_id,
+            subfolder="resized",
+        )
+
+        return original_url, cropped_url, resized_url

visual_rag/indexing/pdf_processor.py

@@ -0,0 +1,324 @@
+"""
+PDF Processor - Convert PDFs to images and extract text.
+
+This module works INDEPENDENTLY of embedding and vector storage.
+Use it if you just need PDF → images conversion.
+
+Features:
+- Batch processing to save memory
+- Text extraction with surrogate character handling
+- Configurable DPI and quality settings
+"""
+
+import gc
+import logging
+import re
+from pathlib import Path
+from typing import Any, Dict, Generator, List, Optional, Tuple
+
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+
+class PDFProcessor:
+    """
+    Process PDFs into images and text for visual retrieval.
+
+    Works independently - no embedding or storage dependencies.
+
+    Args:
+        dpi: DPI for image conversion (higher = better quality)
+        output_format: Image format (RGB, L, etc.)
+        page_batch_size: Pages per batch for memory efficiency
+
+    Example:
+        >>> processor = PDFProcessor(dpi=140)
+        >>>
+        >>> # Convert single PDF
+        >>> images, texts = processor.process_pdf(Path("report.pdf"))
+        >>>
+        >>> # Stream large PDFs
+        >>> for images, texts in processor.stream_pdf(Path("large.pdf"), batch_size=10):
+        ...     # Process each batch
+        ...     pass
+    """
+
+    def __init__(
+        self,
+        dpi: int = 140,
+        output_format: str = "RGB",
+        page_batch_size: int = 50,
+    ):
+        self.dpi = dpi
+        self.output_format = output_format
+        self.page_batch_size = page_batch_size
+
+        # PDF deps are optional: we only require them when calling PDF-specific methods.
+        # This keeps the class usable for helper utilities like `resize_for_colpali()`
+        # even in minimal installs.
+        self._pdf_deps_available = True
+        try:
+            import pdf2image  # noqa: F401
+            import pypdf  # noqa: F401
+        except Exception:
+            self._pdf_deps_available = False
+
+    def _require_pdf_deps(self) -> None:
+        if not self._pdf_deps_available:
+            raise ImportError(
+                "PDF processing requires `pdf2image` and `pypdf`.\n"
+                'Install with: pip install "visual-rag-toolkit[pdf]"'
+            )
+
+    def process_pdf(
+        self,
+        pdf_path: Path,
+        dpi: Optional[int] = None,
+    ) -> Tuple[List[Image.Image], List[str]]:
+        """
+        Convert PDF to images and extract text.
+
+        Args:
+            pdf_path: Path to PDF file
+            dpi: Override default DPI
+
+        Returns:
+            Tuple of (list of images, list of page texts)
+        """
+        self._require_pdf_deps()
+        from pdf2image import convert_from_path
+        from pypdf import PdfReader
+
+        dpi = dpi or self.dpi
+        pdf_path = Path(pdf_path)
+
+        logger.info(f"📄 Processing PDF: {pdf_path.name}")
+
+        # Extract text
+        reader = PdfReader(str(pdf_path))
+        total_pages = len(reader.pages)
+
+        page_texts = []
+        for page in reader.pages:
+            text = page.extract_text() or ""
+            # Handle surrogate characters
+            text = self._sanitize_text(text)
+            page_texts.append(text)
+
+        # Convert to images in batches
+        all_images = []
+        for start_page in range(1, total_pages + 1, self.page_batch_size):
+            end_page = min(start_page + self.page_batch_size - 1, total_pages)
+
+            batch_images = convert_from_path(
+                str(pdf_path),
+                dpi=dpi,
+                fmt=self.output_format.lower(),
+                first_page=start_page,
+                last_page=end_page,
+            )
+
+            all_images.extend(batch_images)
+
+            del batch_images
+            gc.collect()
+
+        assert len(all_images) == len(
+            page_texts
+        ), f"Mismatch: {len(all_images)} images vs {len(page_texts)} texts"
+
+        logger.info(f"✅ Processed {len(all_images)} pages")
+        return all_images, page_texts
+
+    def stream_pdf(
+        self,
+        pdf_path: Path,
+        batch_size: int = 10,
+        dpi: Optional[int] = None,
+    ) -> Generator[Tuple[List[Image.Image], List[str], int], None, None]:
+        """
+        Stream PDF processing for large files.
+
+        Yields batches of (images, texts, start_page) without loading
+        entire PDF into memory.
+
+        Args:
+            pdf_path: Path to PDF file
+            batch_size: Pages per batch
+            dpi: Override default DPI
+
+        Yields:
+            Tuple of (batch_images, batch_texts, start_page_number)
+        """
+        self._require_pdf_deps()
+        from pdf2image import convert_from_path
+        from pypdf import PdfReader
+
+        dpi = dpi or self.dpi
+        pdf_path = Path(pdf_path)
+
+        reader = PdfReader(str(pdf_path))
+        total_pages = len(reader.pages)
+
+        logger.info(f"📄 Streaming PDF: {pdf_path.name} ({total_pages} pages)")
+
+        for start_idx in range(0, total_pages, batch_size):
+            end_idx = min(start_idx + batch_size, total_pages)
+
+            # Extract text for batch
+            batch_texts = []
+            for page_idx in range(start_idx, end_idx):
+                text = reader.pages[page_idx].extract_text() or ""
+                text = self._sanitize_text(text)
+                batch_texts.append(text)
+
+            # Convert images for batch
+            batch_images = convert_from_path(
+                str(pdf_path),
+                dpi=dpi,
+                fmt=self.output_format.lower(),
+                first_page=start_idx + 1,  # 1-indexed
+                last_page=end_idx,
+            )
+
+            yield batch_images, batch_texts, start_idx + 1
+
+            del batch_images
+            gc.collect()
+
+    def get_page_count(self, pdf_path: Path) -> int:
+        """Get number of pages in PDF without loading images."""
+        self._require_pdf_deps()
+        from pypdf import PdfReader
+
+        reader = PdfReader(str(pdf_path))
+        return len(reader.pages)
+
+    def resize_for_colpali(
+        self,
+        image: Image.Image,
+        max_edge: int = 2048,
+        tile_size: int = 512,
+    ) -> Tuple[Image.Image, int, int]:
+        """
+        Resize image following ColPali/Idefics3 processor logic.
+
+        Resizes to fit within tile grid without black padding.
+
+        Args:
+            image: PIL Image
+            max_edge: Maximum edge length
+            tile_size: Size of each tile
+
+        Returns:
+            Tuple of (resized_image, tile_rows, tile_cols)
+        """
+        # Ensure consistent mode for downstream processors (and predictable tests)
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        w, h = image.size
+
+        # Step 1: Resize so longest edge = max_edge
+        if w > h:
+            new_w = max_edge
+            new_h = int(h * (max_edge / w))
+        else:
+            new_h = max_edge
+            new_w = int(w * (max_edge / h))
+
+        # Step 2: Calculate tile grid
+        tile_cols = (new_w + tile_size - 1) // tile_size
+        tile_rows = (new_h + tile_size - 1) // tile_size
+
+        # Step 3: Calculate exact dimensions for tiles
+        final_w = tile_cols * tile_size
+        final_h = tile_rows * tile_size
+
+        # Step 4: Scale to fit within tile grid
+        scale_w = final_w / w
+        scale_h = final_h / h
+        scale = min(scale_w, scale_h)
+
+        scaled_w = int(w * scale)
+        scaled_h = int(h * scale)
+
+        resized = image.resize((scaled_w, scaled_h), Image.LANCZOS)
+
+        # Center on white canvas if needed
+        if scaled_w != final_w or scaled_h != final_h:
+            canvas = Image.new("RGB", (final_w, final_h), (255, 255, 255))
+            offset_x = (final_w - scaled_w) // 2
+            offset_y = (final_h - scaled_h) // 2
+            canvas.paste(resized, (offset_x, offset_y))
+            resized = canvas
+
+        return resized, tile_rows, tile_cols
+
+    def _sanitize_text(self, text: str) -> str:
+        """Remove invalid Unicode characters (surrogates) from text."""
+        if not text:
+            return ""
+
+        # Remove surrogate characters (U+D800-U+DFFF)
+        return text.encode("utf-8", errors="surrogatepass").decode("utf-8", errors="ignore")
+
+    def extract_metadata_from_filename(
+        self,
+        filename: str,
+        mapping: Optional[Dict[str, Dict[str, Any]]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Extract metadata from PDF filename.
+
+        Uses mapping if provided, otherwise falls back to pattern matching.
+
+        Args:
+            filename: PDF filename (with or without .pdf extension)
+            mapping: Optional mapping dict {filename: metadata}
+
+        Returns:
+            Metadata dict with year, source, district, etc.
+        """
+        # Remove extension
+        stem = Path(filename).stem
+        stem_lower = stem.lower().strip()
+
+        # Try mapping first
+        if mapping:
+            if stem_lower in mapping:
+                return mapping[stem_lower].copy()
+
+            # Try without .pdf
+            stem_no_ext = stem_lower.replace(".pdf", "")
+            if stem_no_ext in mapping:
+                return mapping[stem_no_ext].copy()
+
+        # Fallback: pattern matching
+        metadata = {"filename": filename}
+
+        # Extract year
+        year_match = re.search(r"(20\d{2})", stem)
+        if year_match:
+            metadata["year"] = int(year_match.group(1))
+
+        # Detect source type
+        if "consolidated" in stem_lower or ("annual" in stem_lower and "oag" in stem_lower):
+            metadata["source"] = "Consolidated"
+        elif "dlg" in stem_lower or "district local government" in stem_lower:
+            metadata["source"] = "Local Government"
+            # Try to extract district name
+            district_match = re.search(r"([a-z]+)\s+(?:dlg|district local government)", stem_lower)
+            if district_match:
+                metadata["district"] = district_match.group(1).title()
+        elif "hospital" in stem_lower or "referral" in stem_lower:
+            metadata["source"] = "Hospital"
+        elif "ministry" in stem_lower:
+            metadata["source"] = "Ministry"
+        elif "project" in stem_lower:
+            metadata["source"] = "Project"
+        else:
+            metadata["source"] = "Unknown"
+
+        return metadata