File size: 3,229 Bytes
734b5b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from src.agents.annotation_agent import OfflineAnnotationAgent
from src.agents.editing_agent import OfflinePDFEditingAgent
from src.agents.export_agent import OfflineExportAgent
from src.agents.ocr_agent import OfflineOCRAgent
from src.agents.structure_inspector import OfflinePDFStructureInspector
from src.agents.validation_agent import OfflineValidationAgent
from src.core.models import InspectionReport, JobConfig, PDFOperation
@dataclass
class OperationBuckets:
edit: list[PDFOperation]
annotate: list[PDFOperation]
class PDFToolbox:
"""Shared PDF tools used by orchestrators and pipeline definitions."""
EDIT_OPERATION_TYPES = {
"insert_text",
"whiteout_region",
"rotate_page",
"delete_page",
"reorder_pages",
}
ANNOTATION_OPERATION_TYPES = {"highlight", "comment", "rectangle", "redact"}
def __init__(
self,
inspector: OfflinePDFStructureInspector | None = None,
ocr: OfflineOCRAgent | None = None,
editor: OfflinePDFEditingAgent | None = None,
annotator: OfflineAnnotationAgent | None = None,
validator: OfflineValidationAgent | None = None,
exporter: OfflineExportAgent | None = None,
):
self.inspector = inspector or OfflinePDFStructureInspector()
self.ocr_agent = ocr or OfflineOCRAgent()
self.editing_agent = editor or OfflinePDFEditingAgent()
self.annotation_agent = annotator or OfflineAnnotationAgent()
self.validation_agent = validator or OfflineValidationAgent()
self.export_agent = exporter or OfflineExportAgent()
def inspect(self, pdf_path: Path) -> InspectionReport:
return self.inspector.inspect(pdf_path)
def run_ocr(self, input_pdf: Path, output_pdf: Path, language: str = "eng") -> Path:
return self.ocr_agent.run(input_pdf, output_pdf, language)
def edit_pdf(
self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]
) -> Path:
return self.editing_agent.run(input_pdf, output_pdf, operations)
def annotate_pdf(
self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]
) -> Path:
return self.annotation_agent.run(input_pdf, output_pdf, operations)
def validate_pdf(self, input_pdf: Path, output_pdf: Path) -> tuple[Path, dict]:
return self.validation_agent.run(input_pdf, output_pdf)
def export_job(
self, final_pdf: Path, job_config: JobConfig, stages: list, log_file: Path
) -> tuple[Path, Path]:
return self.export_agent.run(final_pdf, job_config, stages, log_file)
def bucket_operations(self, operations: list[PDFOperation]) -> OperationBuckets:
edit_ops = [op for op in operations if op.type in self.EDIT_OPERATION_TYPES]
annotate_ops = [
op for op in operations if op.type in self.ANNOTATION_OPERATION_TYPES
]
return OperationBuckets(edit=edit_ops, annotate=annotate_ops)
def build_default_pdf_toolbox() -> PDFToolbox:
return PDFToolbox()
|