| 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: |
| |
| 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]: |
| |
| 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 |
|
|
