Loaders

from abc import ABC, abstractmethod
from cognity_ai.models import Document

class BaseLoader(ABC):

    @abstractmethod
    def load(self, path: str) -> list[Document]: ...

    @property
    @abstractmethod
    def supported_extensions(self) -> list[str]: ...

from cognity_ai.models import Document, PageInfo, ImageRef

@dataclass
class Document:
    id:          str             # UUID — auto-generated on load
    text:        str             # Full extracted plain text
    metadata:    dict            # source, title, page_count, author, …
    source_path: str             # Absolute path to the original file
    page_map:    list[PageInfo]  # Per-page character offsets + headings
    image_refs:  list[ImageRef]  # Embedded image positions in the text

@dataclass
class PageInfo:
    page_number: int
    char_start:  int
    char_end:    int
    headings:    list[str]

@dataclass
class ImageRef:
    image_id:    str    # Unique ID within the document
    char_offset: int    # Position in text where the image is anchored
    page_number: int
    caption:     str    # OCR-extracted or explicit caption

from cognity_ai.loaders import LoaderFactory

factory = LoaderFactory()

# Auto-detect from extension — returns the appropriate BaseLoader subclass
loader = factory.get_loader("report.pdf")
docs = loader.load("report.pdf")

# Load a directory — recurses and dispatches per-file
all_docs = factory.load_directory("./documents/")

# Register a custom loader for a proprietary format
factory.register(".myext", MyCustomLoader)

from cognity_ai.loaders import PdfLoader, DocxLoader, CsvLoader

# PDF — extract text and images
pdf_loader = PdfLoader(extract_images=True, ocr_fallback=True)
docs = pdf_loader.load("annual_report.pdf")
doc = docs[0]
print(doc.metadata["page_count"])   # → 42
print(len(doc.page_map))            # → 42
print(len(doc.image_refs))          # → number of embedded images

# DOCX — headings and tables
docx_loader = DocxLoader()
docs = docx_loader.load("spec.docx")

# CSV — custom delimiter
csv_loader = CsvLoader(delimiter=";")
docs = csv_loader.load("data.csv")
print(docs[0].metadata["headers"])   # → ["col1", "col2", …]

from cognity_ai.loaders.pdf_utils import (
    extract_tables,   # → list[pd.DataFrame], one per page
    extract_images,   # → list[bytes], raw image bytes per page
    extract_metadata, # → dict  (author, title, created, page_count, …)
    slice_pages,      # → bytes  (sub-PDF from page range)
    merge_pdfs,       # → bytes  (merged single PDF)
    pdf_to_images,    # → list[PIL.Image]  (full-page raster for OCR)
)

from cognity_ai.loaders.pdf_utils import extract_tables, slice_pages, pdf_to_images

# Extract all tables from a financial report
tables = extract_tables("report.pdf")
for df in tables:
    print(df.to_string())

# Slice out the appendix (pages 40–55) as a new PDF
appendix_bytes = slice_pages("report.pdf", start=40, end=55)
with open("appendix.pdf", "wb") as f:
    f.write(appendix_bytes)

# Render pages as images for vision model processing
images = pdf_to_images("scanned_doc.pdf", dpi=300)
for i, img in enumerate(images):
    img.save(f"page_{i+1}.png")

from cognity_ai.loaders import BaseLoader, LoaderFactory
from cognity_ai.models import Document
import uuid

class EpubLoader(BaseLoader):

    @property
    def supported_extensions(self) -> list[str]:
        return [".epub"]

    def load(self, path: str) -> list[Document]:
        # parse the epub, extract text …
        text = _parse_epub(path)
        return [
            Document(
                id=str(uuid.uuid4()),
                text=text,
                metadata={"source": path},
                source_path=path,
                page_map=[],
                image_refs=[],
            )
        ]

# Register globally
factory = LoaderFactory()
factory.register(".epub", EpubLoader)