| """Visual ingest pipeline: PDF -> page images -> ColEmbed embeddings -> store. |
| |
| Pages are embedded in chunks of EMBED_PAGES_PER_CALL so each GPU call stays |
| short; progress events are yielded between chunks so callers can stream them. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| from collections.abc import Iterator |
|
|
| from core.constants import EMBED_PAGES_PER_CALL, RENDER_DPI |
| from core.pdf import page_count, render_pages |
| from core.store import slugify |
| from core.visual_store import VisualStore |
| from models.colembed import ColEmbed |
|
|
|
|
| class VisualIngestPipeline: |
| def __init__(self, embedder: ColEmbed, store: VisualStore): |
| self.embedder = embedder |
| self.store = store |
|
|
| def run( |
| self, pdf_path: str | None, doc_name: str = "", max_pages: int | None = None |
| ) -> Iterator[tuple]: |
| """Index one PDF. Generator yielding ("progress", pages_done, total) |
| after each embedded chunk, then ("done", doc summary dict) last. |
| Re-indexing a manual with the same name overwrites it.""" |
| if not pdf_path: |
| raise ValueError("Please provide a PDF first.") |
| if not pdf_path.lower().endswith(".pdf"): |
| raise ValueError("Only PDFs can be indexed.") |
|
|
| name = doc_name.strip() or ( |
| os.path.splitext(os.path.basename(pdf_path))[0].replace("_", " ") |
| ) |
| doc_id = slugify(name) |
| total = page_count(pdf_path) |
| if max_pages and total > max_pages: |
| raise ValueError( |
| f"This PDF has {total} pages — over the {max_pages}-page cap." |
| ) |
|
|
| writer = self.store.create(doc_id, name, pdf_path, RENDER_DPI, self.embedder.MODEL_ID) |
| try: |
| for start in range(1, total + 1, EMBED_PAGES_PER_CALL): |
| nums = list(range(start, min(start + EMBED_PAGES_PER_CALL, total + 1))) |
| images = render_pages(pdf_path, nums) |
| for num, emb in zip(nums, self.embedder.embed_pages(images)): |
| writer.add_page(num, emb) |
| yield ("progress", nums[-1], total) |
| writer.finalize() |
| except BaseException: |
| writer.abort() |
| raise |
| yield ("done", {"doc_id": doc_id, "name": name, "pages": total}) |
|
|
