from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path from typing import Callable from src.core.models import JobConfig, JobResult, dataclass_to_dict from src.core.toolbox import PDFToolbox from src.services.local_logger import LocalLogger from src.utils.file_utils import ensure_dir, stage_file @dataclass class PipelineContext: job_config: JobConfig job_dir: Path log_file: Path logger: LocalLogger toolbox: PDFToolbox current_pdf: Path stages: list = field(default_factory=list) errors: list = field(default_factory=list) state: dict = field(default_factory=dict) def make_stage_path(self, prefix: str, suffix: str = ".pdf") -> Path: return stage_file(self.job_dir, prefix, suffix=suffix) def finish_success(self, final_pdf: Path, manifest: Path) -> JobResult: self.logger.log("orchestrator", "info", "Job finished", success=True) return JobResult( self.job_config.job_id, True, final_pdf, self.log_file, manifest, self.stages, self.errors, ) def finish_failure(self) -> JobResult: return JobResult( self.job_config.job_id, False, None, self.log_file, self.job_dir / "manifest.json", self.stages, self.errors, ) @dataclass class PipelineStep: name: str action: Callable[[PipelineContext], None] enabled: Callable[[PipelineContext], bool] | None = None def should_run(self, context: PipelineContext) -> bool: if self.enabled is None: return True return self.enabled(context) @dataclass class PDFPipeline: name: str description: str steps: list[PipelineStep] def run(self, context: PipelineContext) -> PipelineContext: context.logger.log( "orchestrator", "info", "Pipeline started", pipeline=self.name, job_id=context.job_config.job_id, ) for step in self.steps: if step.should_run(context): step.action(context) return context def create_pipeline_context( job_config: JobConfig, toolbox: PDFToolbox, work_dir: Path ) -> PipelineContext: job_config.output_dir = ensure_dir(job_config.output_dir) job_dir = ensure_dir(job_config.output_dir / f"job-{job_config.job_id}") log_file = job_dir / "pipeline.jsonl" logger = LocalLogger(log_file) logger.log("orchestrator", "info", "Job started", job_id=job_config.job_id) return PipelineContext( job_config=job_config, job_dir=job_dir, log_file=log_file, logger=logger, toolbox=toolbox, current_pdf=job_config.input_pdf, state={"work_dir": str(work_dir)}, ) def _inspect_step(context: PipelineContext) -> None: report = context.toolbox.inspect(context.current_pdf) buckets = context.toolbox.bucket_operations(context.job_config.operations) context.state["inspection_report"] = report context.state["operation_buckets"] = buckets context.stages.append({"stage": "inspect", "report": dataclass_to_dict(report)}) context.logger.log( "inspect", "info", "Done", ocr_pages=report.ocr_required_pages, edit_operations=len(buckets.edit), annotation_operations=len(buckets.annotate), ) def _needs_ocr(context: PipelineContext) -> bool: report = context.state.get("inspection_report") if report is None: return context.job_config.force_ocr or context.job_config.run_ocr return ( context.job_config.force_ocr or context.job_config.run_ocr or bool(report.ocr_required_pages) ) def _ocr_step(context: PipelineContext) -> None: output = context.make_stage_path("02-ocr") context.current_pdf = context.toolbox.run_ocr( context.current_pdf, output, context.job_config.ocr_language ) context.stages.append({"stage": "ocr", "output": str(context.current_pdf)}) context.logger.log("ocr", "info", "Done", output=str(context.current_pdf)) def _has_edit_ops(context: PipelineContext) -> bool: buckets = context.state.get("operation_buckets") return bool(buckets and buckets.edit) def _edit_step(context: PipelineContext) -> None: buckets = context.state["operation_buckets"] output = context.make_stage_path("03-edit") context.current_pdf = context.toolbox.edit_pdf( context.current_pdf, output, buckets.edit ) context.stages.append( { "stage": "edit", "output": str(context.current_pdf), "count": len(buckets.edit), } ) context.logger.log("edit", "info", "Done", count=len(buckets.edit)) def _has_annotation_ops(context: PipelineContext) -> bool: buckets = context.state.get("operation_buckets") return bool(buckets and buckets.annotate) def _annotate_step(context: PipelineContext) -> None: buckets = context.state["operation_buckets"] output = context.make_stage_path("04-annotate") context.current_pdf = context.toolbox.annotate_pdf( context.current_pdf, output, buckets.annotate ) context.stages.append( { "stage": "annotate", "output": str(context.current_pdf), "count": len(buckets.annotate), } ) context.logger.log("annotate", "info", "Done", count=len(buckets.annotate)) def _validate_step(context: PipelineContext) -> None: output = context.make_stage_path("05-validated") context.current_pdf, report = context.toolbox.validate_pdf( context.current_pdf, output ) context.state["validation_report"] = report context.stages.append( {"stage": "validate", "output": str(context.current_pdf), "report": report} ) context.logger.log("validate", "info", "Done", **report) def _export_step(context: PipelineContext) -> None: final_pdf, manifest = context.toolbox.export_job( context.current_pdf, context.job_config, context.stages, context.log_file, ) context.state["final_pdf"] = final_pdf context.state["manifest"] = manifest context.stages.append( {"stage": "export", "output": str(final_pdf), "manifest": str(manifest)} ) context.logger.log("export", "info", "Done", final=str(final_pdf)) def build_pipeline_catalog(toolbox: PDFToolbox) -> dict[str, PDFPipeline]: inspect = PipelineStep("inspect", _inspect_step) ocr = PipelineStep("ocr", _ocr_step, enabled=_needs_ocr) edit = PipelineStep("edit", _edit_step, enabled=_has_edit_ops) annotate = PipelineStep("annotate", _annotate_step, enabled=_has_annotation_ops) validate = PipelineStep("validate", _validate_step) export = PipelineStep("export", _export_step) return { "full_document": PDFPipeline( name="full_document", description="Inspects the document, runs OCR if needed, applies edits and annotations, validates, then exports.", steps=[inspect, ocr, edit, annotate, validate, export], ), "text_edit": PDFPipeline( name="text_edit", description="Runs the edit-oriented document pipeline without annotation steps.", steps=[inspect, ocr, edit, validate, export], ), "annotation_review": PDFPipeline( name="annotation_review", description="Applies annotation and review operations on top of the current PDF, then validates and exports.", steps=[inspect, annotate, validate, export], ), "page_restructure": PDFPipeline( name="page_restructure", description="Handles page-level editing operations like rotate, delete, and reorder before validation and export.", steps=[inspect, edit, validate, export], ), "ocr_preflight": PDFPipeline( name="ocr_preflight", description="Runs inspection and OCR preparation only, then validates and exports an OCR-ready PDF.", steps=[inspect, ocr, validate, export], ), } def list_available_pipeline_names(toolbox: PDFToolbox | None = None) -> list[str]: resolved_toolbox = toolbox or PDFToolbox() return list(build_pipeline_catalog(resolved_toolbox).keys())