src/ingestion.py

import json
import os
import re
from dataclasses import dataclass, asdict
from typing import List, Optional, Tuple

from pypdf import PdfReader
import pdfplumber


@dataclass
class Chunk:
    content: str
    source: str
    page: Optional[int] = None
    chapter: Optional[str] = None


def normalize_text(text: str) -> str:
    # Remove control characters, normalize whitespace
    text = text.encode("utf-8", errors="ignore").decode("utf-8", errors="ignore")
    text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F]", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def extract_text_pypdf(file_path: str) -> List[Tuple[int, str]]:
    pages: List[Tuple[int, str]] = []
    reader = PdfReader(file_path)
    for i, page in enumerate(reader.pages, start=1):
        text = page.extract_text() or ""
        pages.append((i, text))
    return pages


def extract_text_pdfplumber(file_path: str) -> List[Tuple[int, str]]:
    pages: List[Tuple[int, str]] = []
    with pdfplumber.open(file_path) as pdf:
        for i, page in enumerate(pdf.pages, start=1):
            text = page.extract_text() or ""
            pages.append((i, text))
    return pages


def chunk_text(text: str, chunk_size: int = 4000, overlap: int = 400) -> List[str]:
    # Simple recursive-like character splitter
    if chunk_size <= 0:
        return [text]
    chunks: List[str] = []
    start = 0
    n = len(text)
    while start < n:
        end = min(start + chunk_size, n)
        chunk = text[start:end]
        chunks.append(chunk)
        if end == n:
            break
        start = max(end - overlap, start + 1)
    return chunks


def make_chunks(pages: List[Tuple[int, str]], source: str, chunk_size: int, overlap: int) -> List[Chunk]:
    all_text = []
    for page_num, raw in pages:
        norm = normalize_text(raw)
        if norm:
            all_text.append((page_num, norm))
    merged = "\n".join(t for _, t in all_text)
    pieces = chunk_text(merged, chunk_size=chunk_size, overlap=overlap)
    chunks = [Chunk(content=p, source=source) for p in pieces if p.strip()]
    return chunks


def persist_chunks(session_id: str, chunks: List[Chunk]) -> str:
    base = os.path.join(".tmp", "ingestion")
    os.makedirs(base, exist_ok=True)
    path = os.path.join(base, f"{session_id}.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump([asdict(c) for c in chunks], f, ensure_ascii=False, indent=2)
    return path


def ingest_pdf(
    file_path: str,
    session_id: str,
    parser: str = "pypdf",
    chunk_size: int = 4000,
    overlap: int = 400,
) -> Tuple[List[Chunk], str]:
    if parser == "pypdf":
        pages = extract_text_pypdf(file_path)
    elif parser == "pdfplumber":
        pages = extract_text_pdfplumber(file_path)
    else:
        raise ValueError("Unknown parser: " + parser)

    chunks = make_chunks(pages, source=os.path.basename(file_path), chunk_size=chunk_size, overlap=overlap)
    out = persist_chunks(session_id, chunks)
    return chunks, out