utils/pdf_processor.py

"""utils/pdf_processor.py — PDF ingestion pipeline for the HF Space.

Standalone module mirroring src/rag_system/components/layout_parser.py logic:
PDF -> page images + raw text -> semantic chunk -> ready for embedding.
"""
from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional, Tuple


@dataclass
class DocumentChunk:
    text: str
    page_number: int
    chunk_index: int
    source_filename: str
    chunk_type: str = "text"
    table_data: Optional[str] = None
    metadata: dict = field(default_factory=dict)

    @property
    def chunk_id(self) -> str:
        content = f"{self.source_filename}:{self.page_number}:{self.chunk_index}"
        return hashlib.sha256(content.encode()).hexdigest()[:12]


@dataclass
class IngestResult:
    filename: str
    num_pages: int
    num_chunks: int
    num_tables: int
    num_charts: int
    chunks: List[DocumentChunk]
    page_images: List[object]
    processing_steps: List[str]


def extract_text_and_tables(pdf_path: str) -> Tuple[List[dict], List[str]]:
    steps = []
    pages = []
    try:
        import pdfplumber
        steps.append("Opened PDF with pdfplumber")
        with pdfplumber.open(pdf_path) as pdf:
            steps.append(f"Found {len(pdf.pages)} pages")
            for i, page in enumerate(pdf.pages, 1):
                text = page.extract_text() or ""
                tables = page.extract_tables() or []
                pages.append({"page_num": i, "text": text.strip(), "tables": tables})
            steps.append(f"Extracted text from {len(pages)} pages")
            total_tables = sum(len(p["tables"]) for p in pages)
            if total_tables:
                steps.append(f"Found {total_tables} tables across all pages")
    except ImportError:
        steps.append("pdfplumber not available - trying PyPDF2 fallback")
        try:
            import PyPDF2
            with open(pdf_path, "rb") as f:
                reader = PyPDF2.PdfReader(f)
                for i, page in enumerate(reader.pages, 1):
                    text = page.extract_text() or ""
                    pages.append({"page_num": i, "text": text.strip(), "tables": []})
            steps.append(f"Extracted text (no table detection) from {len(pages)} pages")
        except Exception as e:
            steps.append(f"PDF extraction failed: {str(e)[:80]}")
    except Exception as e:
        steps.append(f"Extraction error: {str(e)[:80]}")
    return pages, steps


def render_page_images(pdf_path: str, max_pages: int = 8) -> Tuple[List[object], List[str]]:
    steps = []
    images = []
    try:
        from pdf2image import convert_from_path
        images = convert_from_path(pdf_path, dpi=150, first_page=1, last_page=max_pages)
        steps.append(f"Rendered {len(images)} page images at 150 DPI for vision analysis")
    except ImportError:
        steps.append("pdf2image not available - skipping page rendering")
    except Exception as e:
        steps.append(f"Page rendering failed: {str(e)[:80]}")
    return images, steps


def semantic_chunk_text(
    text: str,
    page_num: int,
    source: str,
    chunk_start_index: int = 0,
    max_chars: int = 800,
    overlap_chars: int = 100,
) -> List[DocumentChunk]:
    if not text.strip():
        return []
    paragraphs = [p.strip() for p in re.split(r"\n{2,}", text) if p.strip()]
    chunks: List[DocumentChunk] = []
    buffer = ""
    chunk_idx = chunk_start_index
    for para in paragraphs:
        if len(buffer) + len(para) < max_chars:
            buffer = (buffer + "\n\n" + para).strip() if buffer else para
        else:
            if buffer:
                chunks.append(DocumentChunk(
                    text=buffer, page_number=page_num, chunk_index=chunk_idx,
                    source_filename=source, chunk_type="text",
                ))
                chunk_idx += 1
                buffer = buffer[-overlap_chars:] + "\n\n" + para if len(buffer) > overlap_chars else para
            else:
                buffer = para
    if buffer.strip():
        chunks.append(DocumentChunk(
            text=buffer.strip(), page_number=page_num, chunk_index=chunk_idx,
            source_filename=source, chunk_type="text",
        ))
    return chunks


def format_table_as_text(table: List[List[str]]) -> str:
    if not table:
        return ""
    rows = []
    for row in table:
        cells = [str(c or "").strip() for c in row]
        rows.append(" | ".join(cells))
    return "\n".join(rows)


def chunk_tables(
    tables: List[List[List[str]]],
    page_num: int,
    source: str,
    chunk_start_index: int = 0,
) -> List[DocumentChunk]:
    chunks = []
    for t_idx, table in enumerate(tables):
        table_text = format_table_as_text(table)
        if table_text.strip():
            chunks.append(DocumentChunk(
                text=f"[TABLE on page {page_num}]\n{table_text}",
                page_number=page_num, chunk_index=chunk_start_index + t_idx,
                source_filename=source, chunk_type="table",
                table_data=table_text,
            ))
    return chunks


def ingest_pdf(pdf_path: str, process_vision: bool = True, vision_fn=None) -> IngestResult:
    filename = Path(pdf_path).name
    steps: List[str] = [f"Processing: {filename}"]
    all_chunks: List[DocumentChunk] = []
    page_images = []

    pages, extract_steps = extract_text_and_tables(pdf_path)
    steps.extend(extract_steps)

    chunk_counter = 0
    num_tables = 0
    for page in pages:
        text_chunks = semantic_chunk_text(page["text"], page["page_num"], filename, chunk_counter)
        all_chunks.extend(text_chunks)
        chunk_counter += len(text_chunks)
        table_chunks = chunk_tables(page["tables"], page["page_num"], filename, chunk_counter)
        all_chunks.extend(table_chunks)
        num_tables += len(table_chunks)
        chunk_counter += len(table_chunks)

    steps.append(f"Chunked into {len(all_chunks)} semantic chunks ({num_tables} tables)")

    num_charts = 0
    if process_vision and vision_fn is not None:
        imgs, img_steps = render_page_images(pdf_path)
        steps.extend(img_steps)
        page_images = imgs
        steps.append("Running vision LLM on page images...")
        for img_idx, img in enumerate(imgs[:6]):
            try:
                description = vision_fn(img)
                if description and len(description) > 30:
                    all_chunks.append(DocumentChunk(
                        text=f"[VISUAL CONTENT - page {img_idx + 1}]\n{description}",
                        page_number=img_idx + 1, chunk_index=chunk_counter,
                        source_filename=filename, chunk_type="chart_description",
                    ))
                    chunk_counter += 1
                    num_charts += 1
            except Exception as e:
                steps.append(f"Vision failed for page {img_idx + 1}: {str(e)[:60]}")
        if num_charts:
            steps.append(f"Generated {num_charts} visual descriptions from page images")
    elif process_vision:
        steps.append("Vision skipped - no vision model provided")

    steps.append(f"Ingestion complete: {len(all_chunks)} total chunks ready for retrieval")

    return IngestResult(
        filename=filename, num_pages=len(pages), num_chunks=len(all_chunks),
        num_tables=num_tables, num_charts=num_charts, chunks=all_chunks,
        page_images=page_images, processing_steps=steps,
    )