pdf-reader-0jsfjf 0.1.0__py3-none-any.whl

pdf_reader/__init__.py

@@ -0,0 +1,45 @@
+"""
+PDF Reader - PDF indexing, search, and retrieval for AI applications.
+
+This package provides:
+- PDF/PDFPage: Data models for PDF documents and pages
+- PDFIndex: Hybrid index combining Tantivy search with SQLite storage
+- PDFService: Multi-index service layer for business logic
+- Server implementations: REST API, MCP, and lm_deluge tools
+- Formatting helpers: format_toc, format_ocr for LLM consumption
+"""
+
+from .index import PDFIndex
+from .pdf import PDF, PDFPage, format_toc, format_ocr
+from .service import PDFService, PDFServiceConfig
+from .ingest import ingest_pages_llm, iter_ingest_pages_llm
+from .core import (
+    PDFIndexConfig,
+    SearchRequest,
+    SearchResponse,
+    ReadRequest,
+    ReadResponse,
+)
+
+__all__ = [
+    # Data models
+    "PDF",
+    "PDFPage",
+    # Index
+    "PDFIndex",
+    # Service
+    "PDFService",
+    "PDFServiceConfig",
+    # Config and API types
+    "PDFIndexConfig",
+    "SearchRequest",
+    "SearchResponse",
+    "ReadRequest",
+    "ReadResponse",
+    # Formatting helpers
+    "format_toc",
+    "format_ocr",
+    # Streaming OCR
+    "ingest_pages_llm",
+    "iter_ingest_pages_llm",
+]

pdf_reader/core.py

@@ -0,0 +1,150 @@
+"""
+Core types, models, and utilities shared across all PDF service implementations.
+
+This module provides:
+- PDFIndexConfig dataclass for index configuration
+- Request/Response Pydantic models for the unified API
+- Utility functions for data processing
+"""
+
+from dataclasses import dataclass
+
+from pydantic import BaseModel, Field
+
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+
+@dataclass
+class PDFIndexConfig:
+    """Configuration for a single PDF index."""
+
+    name: str
+    description: str = ""
+    root_path: str | None = None  # Directory with search/, pages.db, manifest.json
+
+
+# =============================================================================
+# Request/Response Models (Unified API)
+# =============================================================================
+
+
+class SearchRequest(BaseModel):
+    """Request model for search endpoint."""
+
+    index: str = Field(..., description="Name of the index to search")
+    queries: list[str] = Field(..., description="List of search query strings")
+    limit: int = Field(10, description="Maximum number of results to return")
+    pdf_filter: str | None = Field(
+        None, description="Optional PDF ID to filter results to a single document"
+    )
+    include_images: bool = Field(
+        False, description="If true, include page images in results"
+    )
+
+
+class PageResult(BaseModel):
+    """A single page result from search."""
+
+    id: str = Field(..., description="Page ID in format 'pdf_id:page_index'")
+    pdf_id: str = Field(..., description="ID of the parent PDF")
+    page_index: int = Field(..., description="Zero-based page number")
+    text: str = Field("", description="OCR-extracted text")
+    summary: str | None = Field(None, description="AI-generated page summary")
+    caption: str | None = Field(
+        None, description="AI-generated caption for visual pages"
+    )
+    image: str | None = Field(None, description="Base64-encoded JPEG image")
+    preview: str | None = Field(None, description="Text preview for display")
+
+
+class SearchResponse(BaseModel):
+    """Response model for search endpoint."""
+
+    index: str
+    queries: list[str]
+    results: list[PageResult]
+    total: int
+
+
+class ReadRequest(BaseModel):
+    """Request model for read endpoint."""
+
+    index: str = Field(..., description="Name of the index to read from")
+    pdf_id: str = Field(..., description="ID of the PDF to read")
+    pages: list[int] | str = Field(
+        ...,
+        description="List of page indices to read, or a range string like '1-5' or '-1' for all",
+    )
+    include_images: bool = Field(True, description="If true, include page images")
+
+
+class ReadResponse(BaseModel):
+    """Response model for read endpoint."""
+
+    index: str
+    pdf_id: str
+    pages: list[PageResult]
+
+
+class PDFInfo(BaseModel):
+    """Information about a single PDF in an index."""
+
+    id: str
+    filename: str
+    title: str | None = None
+    summary: str | None = None
+    num_pages: int
+
+
+class ListPDFsResponse(BaseModel):
+    """Response model for listing PDFs in an index."""
+
+    index: str
+    pdfs: list[PDFInfo]
+
+
+class IndexInfo(BaseModel):
+    """Information about a single index."""
+
+    name: str
+    description: str = ""
+    pdf_count: int = 0
+
+
+class IndicesResponse(BaseModel):
+    """Response model for listing indices."""
+
+    indices: list[IndexInfo]
+
+
+class HealthResponse(BaseModel):
+    """Response model for health check."""
+
+    status: str
+    indices: list[str]
+    server: str | None = None
+
+
+# =============================================================================
+# Exports
+# =============================================================================
+
+__all__ = [
+    # Configuration
+    "PDFIndexConfig",
+    # Request models
+    "SearchRequest",
+    "ReadRequest",
+    # Response models
+    "PageResult",
+    "SearchResponse",
+    "ReadResponse",
+    "PDFInfo",
+    "ListPDFsResponse",
+    "IndexInfo",
+    "IndicesResponse",
+    "HealthResponse",
+]

pdf_reader/index/__init__.py

@@ -0,0 +1,370 @@
+from __future__ import annotations
+
+import asyncio
+import glob
+import json
+import shutil
+import tempfile
+import zipfile
+from pathlib import Path
+from typing import List, Literal, TypedDict
+
+from pdf_reader.util import fetch_all_pdfs_concurrently, generate_pdf_hash, parse_range
+
+from ..ingest import enrich, ingest_many
+from ..pdf import PDF, PDFPage
+from .page_store import PageStore
+from .tantivy_ext import PageSearch
+
+
+class File(TypedDict):
+    id: str
+    filename: str
+    presigned_url: str | None
+    bytes: bytes | None
+
+
+class PDFIndex:
+    """
+    Hybrid index: Tantivy for retrieval; SQLite for random page access.
+    Directory layout:
+
+        <root>/
+            search/     (tantivy index)
+            pages.db    (sqlite blob store)
+            manifest.json
+    """
+
+    def __init__(self, root: str, search: PageSearch, store: PageStore):
+        self.root = Path(root)
+        self.search = search
+        self.store = store
+        with open(self.root / "manifest.json") as f:
+            self.manifest = json.load(f)  # maps pdf_id -> {title, num_pages}
+
+    # ---------- builders ----------
+    @classmethod
+    async def from_folder(
+        cls,
+        folder: str | Path,
+        root: str | Path,
+        do_enrich: bool = True,
+        ocr_backend: Literal["tesseract", "marker", "llm"] = "llm",
+        backend_llm: str | None = "gemini-2.5-flash",
+        cache_path: str | None = None,
+        max_concurrent_requests: int = 225,
+        keep_raw_bytes: bool = False,
+    ) -> "PDFIndex":
+        pdf_paths = glob.glob(f"{folder}/*.pdf")
+        filenames = [x.split("/")[-1] for x in pdf_paths]
+
+        return await PDFIndex.from_pdfs(
+            pdf_paths,
+            root,
+            filenames,  # type: ignore
+            do_enrich,
+            ocr_backend,
+            backend_llm,
+            cache_path,
+            max_concurrent_requests,
+            keep_raw_bytes=keep_raw_bytes,
+        )
+
+    @classmethod
+    async def from_pdfs(
+        cls,
+        pdf_paths: List[str],
+        root: str | Path,
+        filenames: list[str | None] | None = None,
+        do_enrich: bool = True,
+        ocr_backend: Literal["tesseract", "marker", "llm"] = "llm",
+        backend_llm: str | None = "gemini-2.5-flash",
+        cache_path: str | None = None,
+        max_concurrent_requests: int = 225,
+        keep_raw_bytes: bool = False,
+    ) -> "PDFIndex":
+        root = Path(root)
+        if root.exists():
+            print(f"⚠️  Removing existing index dir {root}")
+            shutil.rmtree(root)
+        (root / "search").mkdir(parents=True, exist_ok=True)
+
+        # Ingest all PDFs
+        all_pages: List[PDFPage] = []
+        manifest = {}
+        tasks = []
+        if not filenames:
+            filenames = [None for _ in pdf_paths]
+
+        print("ingesting pdfs...")
+        pdfs: list[PDF] = await ingest_many(
+            pdf_paths,
+            filenames,
+            backend=ocr_backend,
+            backend_llm=backend_llm,
+            max_concurrent_requests=max_concurrent_requests,
+            keep_raw_bytes=keep_raw_bytes,
+        )
+
+        # if enriching, enrich the pdfs
+        print("enriching pdfs...")
+        if do_enrich:
+            tasks = []
+            for pdf in pdfs:
+                tasks.append(asyncio.create_task(enrich(pdf, cache_path=cache_path)))
+
+            await asyncio.gather(*tasks)
+
+        for pdf in pdfs:
+            manifest[pdf.id] = {
+                "id": pdf.id,
+                "pdf_id": pdf.id,
+                "title": pdf.title,
+                "filename": pdf.filename,
+                "summary": pdf.summary,
+                "pages": pdf.num_pages,
+            }
+            all_pages.extend(pdf.pages)
+
+        # Build search index
+        search = PageSearch(str(root / "search"))
+        search.build(pdfs)
+
+        # Build page store (with raw PDF bytes if requested)
+        store = PageStore(str(root / "pages.db"))
+        store.build(all_pages, pdfs=pdfs if keep_raw_bytes else None)
+
+        # Persist manifest
+        with open(root / "manifest.json", "w") as f:
+            json.dump(manifest, f, indent=2)
+
+        return cls(str(root), search, store)
+
+    @classmethod
+    def from_pdf_objects(
+        cls,
+        pdfs: List["PDF"],
+        root: str | Path,
+        keep_raw_bytes: bool = False,
+    ) -> "PDFIndex":
+        """
+        Build an index from pre-constructed PDF objects.
+
+        This is useful when PDFs have already been OCR'd and loaded
+        (e.g., from S3 flat files) and you just need to build the index.
+
+        Args:
+            pdfs: List of PDF objects with pages already populated
+            root: Directory to store the index
+            keep_raw_bytes: Whether to store raw PDF bytes (if available on PDF objects)
+
+        Returns:
+            PDFIndex ready for searching
+        """
+        root = Path(root)
+        if root.exists():
+            shutil.rmtree(root)
+        (root / "search").mkdir(parents=True, exist_ok=True)
+
+        # Build manifest and collect pages
+        manifest = {}
+        all_pages: List[PDFPage] = []
+
+        for pdf in pdfs:
+            manifest[pdf.id] = {
+                "id": pdf.id,
+                "pdf_id": pdf.id,
+                "title": pdf.title,
+                "filename": pdf.filename,
+                "summary": pdf.summary,
+                "pages": len(pdf.pages),
+            }
+            all_pages.extend(pdf.pages)
+
+        # Build search index
+        search = PageSearch(str(root / "search"))
+        search.build(pdfs)
+
+        # Build page store
+        store = PageStore(str(root / "pages.db"))
+        store.build(
+            all_pages,
+            pdfs=pdfs if keep_raw_bytes and any(p.raw_bytes for p in pdfs) else None,
+        )
+
+        # Persist manifest
+        with open(root / "manifest.json", "w") as f:
+            json.dump(manifest, f, indent=2)
+
+        return cls(str(root), search, store)
+
+    @classmethod
+    def load(cls, root: str | Path) -> "PDFIndex":
+        root = Path(root)
+        if not (root / "manifest.json").exists():
+            raise FileNotFoundError("Not a valid PDFIndex directory")
+        search = PageSearch.open(str(root / "search"))
+        store = PageStore.open(str(root / "pages.db"))
+        return cls(str(root), search, store)
+
+    def save(self, *_):
+        """No‑op – everything is already on disk."""
+        pass
+
+    def save_to_file(self, filepath: str | Path) -> None:
+        """Save the entire PDFIndex to a single ZIP file."""
+        filepath = Path(filepath)
+        with zipfile.ZipFile(filepath, "w", zipfile.ZIP_DEFLATED) as zipf:
+            for file_path in self.root.rglob("*"):
+                if file_path.is_file():
+                    arcname = file_path.relative_to(self.root)
+                    zipf.write(file_path, arcname)
+
+    @classmethod
+    def load_from_file(
+        cls, filepath: str | Path, extract_to: str | Path | None = None
+    ) -> "PDFIndex":
+        """Load a PDFIndex from a single ZIP file."""
+        filepath = Path(filepath)
+
+        if extract_to is None:
+            extract_to = tempfile.mkdtemp(prefix="pdfindex_")
+        else:
+            extract_to = Path(extract_to)
+
+        with zipfile.ZipFile(filepath, "r") as zipf:
+            zipf.extractall(extract_to)
+
+        return cls.load(extract_to)
+
+    # ---------- API ----------
+    def search_pages(
+        self, queries: List[str], pdf_id: str | None = None, limit: int = 10
+    ) -> List[PDFPage]:
+        """Keyword search across all PDFs (or a single PDF)."""
+        ids = self.search.search(queries, pdf_filter=pdf_id, limit=limit)
+        return [self.store.fetch_page(pid) for pid in ids]
+
+    def read(self, pdf_id: str, pages: List[int] | str) -> List[PDFPage]:
+        if isinstance(pages, str):
+            pages = parse_range(pages)
+        return self.store.fetch_pages(pdf_id, pages)
+
+    def list_pdfs(self) -> List[dict]:
+        """List all PDFs in the index with their metadata."""
+        # could just do return self.manifest.values()
+        return [metadata for pdf_id, metadata in self.manifest.items()]
+
+    def get_pdfs(self, include_pages: bool = True) -> List[PDF]:
+        """Get all PDFs as PDF objects.
+
+        Args:
+            include_pages: If True, include all pages. If False, return PDFs with empty page lists.
+
+        Returns:
+            List of PDF objects with metadata and optionally pages.
+        """
+        pdfs = []
+        for pdf_id, meta in self.manifest.items():
+            if include_pages:
+                pages = self.store.fetch_pages(pdf_id, [-1])  # -1 means all pages
+            else:
+                pages = []
+
+            pdf = PDF(
+                id=pdf_id,
+                filename=meta.get("filename", ""),
+                title=meta.get("title"),
+                summary=meta.get("summary"),
+                pages=pages,
+                raw_bytes=self.store.fetch_raw_bytes(pdf_id),
+            )
+            pdfs.append(pdf)
+        return pdfs
+
+    # ---------- raw PDF access ----------
+    def get_raw_bytes(self, pdf_id: str) -> bytes | None:
+        """Get raw PDF bytes for a specific PDF. Returns None if not stored."""
+        return self.store.fetch_raw_bytes(pdf_id)
+
+    def get_all_raw_bytes(self) -> dict[str, bytes]:
+        """Get all raw PDF bytes. Returns dict of pdf_id -> bytes."""
+        return self.store.fetch_all_raw_bytes()
+
+    def has_raw_bytes(self) -> bool:
+        """Check if this index has raw PDF bytes stored."""
+        return self.store.has_pdfs_table()
+
+
+async def get_or_create_pdf_index(
+    files: list[File],
+    cache_path: str,
+    llm_cache_path: str,
+    max_concurrent_requests: int = 225,
+) -> tuple[PDFIndex, str]:
+    """Get cached PDF index or create new one if cache miss."""
+
+    presigned_urls = [
+        file["presigned_url"] for file in files if file.get("presigned_url")
+    ]
+    bytes_pdf_contents = [file["bytes"] for file in files if file.get("bytes")]
+
+    if presigned_urls:
+        print("Fetching PDFs concurrently...")
+        url_pdf_contents = await fetch_all_pdfs_concurrently(presigned_urls)  # type: ignore
+    else:
+        url_pdf_contents = []
+
+    pdf_contents = url_pdf_contents + bytes_pdf_contents
+    filenames = [f["filename"] for f in files]
+
+    # Generate order-invariant hash
+    content_hash = generate_pdf_hash(pdf_contents)  # type: ignore
+    print(f"Generated content hash: {content_hash}")
+
+    # Check if cached index exists
+    cache_dir = Path(cache_path)
+    cache_dir.mkdir(exist_ok=True)
+    cache_file = cache_dir / f"{content_hash}.pkl"
+
+    if cache_file.exists():
+        print(f"Loading cached index from {cache_file}")
+        return PDFIndex.load_from_file(str(cache_file)), content_hash
+    else:
+        print(f"Cache miss. Creating new index and saving to {cache_file}")
+
+        # Save PDFs to temporary files
+        temp_files = []
+        try:
+            for i, content in enumerate(pdf_contents):
+                if not content:
+                    continue
+                temp_file = tempfile.NamedTemporaryFile(
+                    delete=False, suffix=f"_{i}.pdf"
+                )
+                temp_file.write(content)
+                temp_file.close()
+                temp_files.append(temp_file.name)
+
+            # Create new index from local files
+            index = await PDFIndex.from_pdfs(
+                temp_files,
+                filenames=filenames,  # type: ignore
+                root="./pdf-index",
+                ocr_backend="llm",
+                backend_llm="gemini-2.5-flash",
+                cache_path=llm_cache_path,
+                max_concurrent_requests=max_concurrent_requests,
+            )
+
+            # Save to cache
+            index.save_to_file(str(cache_file))
+            return index, content_hash
+
+        finally:
+            # Clean up temporary files
+            for temp_file in temp_files:
+                try:
+                    Path(temp_file).unlink()
+                except FileNotFoundError:
+                    pass