| 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()
|
|
|
