| from __future__ import annotations |
| from pathlib import Path |
| import fitz |
| from pypdf import PdfReader, PdfWriter |
| from src.core.models import PDFOperation |
|
|
| class OfflinePDFEditingAgent: |
| name = "editing_agent" |
|
|
| def run(self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]) -> Path: |
| doc = fitz.open(input_pdf) |
| page_ops = [op for op in operations if op.type in ("insert_text","whiteout_region","rotate_page")] |
| delete_ops = [op for op in operations if op.type == "delete_page"] |
| reorder_ops = [op for op in operations if op.type == "reorder_pages"] |
|
|
| for op in page_ops: |
| if op.page is None or op.page >= len(doc): continue |
| page = doc[op.page] |
| if op.type == "insert_text" and op.position and op.text: |
| color = tuple(op.color) if op.color else (0, 0, 0) |
| page.insert_text(op.position, op.text, fontsize=12, color=color, overlay=True) |
| elif op.type == "whiteout_region" and op.rect: |
| page.draw_rect(fitz.Rect(op.rect), color=(1,1,1), fill=(1,1,1), overlay=True) |
| if op.text: |
| x0, y0 = op.rect[0], op.rect[1] |
| page.insert_text((x0+4, y0+14), op.text, fontsize=12, color=(0,0,0), overlay=True) |
| elif op.type == "rotate_page" and op.rotation is not None: |
| page.set_rotation(op.rotation) |
|
|
| stage = output_pdf.with_name(output_pdf.stem + "-pre.pdf") |
| stage.parent.mkdir(parents=True, exist_ok=True) |
| doc.save(stage, garbage=4, deflate=True, clean=True) |
| doc.close() |
|
|
| reader = PdfReader(str(stage)) |
| writer = PdfWriter() |
| delete_indices = {op.page for op in delete_ops if op.page is not None} |
| order = reorder_ops[-1].pages if reorder_ops else None |
| indices = [i for i in range(len(reader.pages)) if i not in delete_indices] |
| if order: |
| indices = [i for i in order if i not in delete_indices] |
| for idx in indices: |
| writer.add_page(reader.pages[idx]) |
| output_pdf.parent.mkdir(parents=True, exist_ok=True) |
| with output_pdf.open("wb") as f: |
| writer.write(f) |
| stage.unlink(missing_ok=True) |
| return output_pdf |
|
|
