105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.logging_config import get_logger
|
|
from app.models.document import Document
|
|
from app.repositories.document_repository import DocumentRepository
|
|
from app.services.layout_service import LayoutService
|
|
from app.services.ocr_service import OCRService
|
|
from app.services.pdf_service import NativePDFService
|
|
from app.services.template_service import TemplateService
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class DocumentProcessingService:
|
|
"""Orchestrates the full document processing pipeline."""
|
|
|
|
def __init__(self, db: Session) -> None:
|
|
self.db = db
|
|
self.doc_repo = DocumentRepository(db)
|
|
self.pdf_service = NativePDFService(db)
|
|
self.ocr_service = OCRService(db)
|
|
self.layout_service = LayoutService(db)
|
|
self.template_service = TemplateService(db)
|
|
|
|
def process_document(self, document_id: str | uuid.UUID) -> Document:
|
|
"""Process a document through the full pipeline."""
|
|
if isinstance(document_id, str):
|
|
document_id = uuid.UUID(document_id)
|
|
|
|
document = self.doc_repo.get_by_id(document_id)
|
|
if not document:
|
|
raise ValueError(f"Document '{document_id}' not found")
|
|
|
|
logger.info(
|
|
"processing_started",
|
|
document_id=str(document_id),
|
|
content_type=document.content_type,
|
|
)
|
|
|
|
# Update status to processing
|
|
self.doc_repo.update_status(document_id, "processing")
|
|
self.db.commit()
|
|
|
|
try:
|
|
# Step 1: Extract content based on document type
|
|
if document.content_type == "application/pdf":
|
|
document = self._process_pdf(document)
|
|
else:
|
|
document = self._process_image(document)
|
|
|
|
# Step 2: Analyze layout
|
|
layout_results = self.layout_service.analyze_document_layout(document)
|
|
metadata = document.document_metadata or {}
|
|
metadata["layout"] = layout_results
|
|
document.document_metadata = metadata
|
|
|
|
# Step 3: Update document status
|
|
self.doc_repo.update_status(document_id, "completed")
|
|
self.db.commit()
|
|
|
|
logger.info(
|
|
"processing_completed",
|
|
document_id=str(document_id),
|
|
pages=document.page_count,
|
|
)
|
|
|
|
return document
|
|
|
|
except Exception as e:
|
|
logger.exception(
|
|
"processing_failed",
|
|
document_id=str(document_id),
|
|
error=str(e),
|
|
)
|
|
self.doc_repo.update_status(document_id, "failed", error_message=str(e))
|
|
self.db.commit()
|
|
raise
|
|
|
|
def _process_pdf(self, document: Document) -> Document:
|
|
"""Process a PDF document - either native or scanned."""
|
|
# First, try native PDF extraction
|
|
document = self.pdf_service.process_pdf(document)
|
|
self.db.flush()
|
|
|
|
# If scanned, also run OCR
|
|
if document.is_scanned:
|
|
logger.info(
|
|
"scanned_pdf_detected",
|
|
document_id=str(document.id),
|
|
)
|
|
document = self.ocr_service.process_scanned_pdf(document)
|
|
self.db.flush()
|
|
|
|
return document
|
|
|
|
def _process_image(self, document: Document) -> Document:
|
|
"""Process an image document with OCR."""
|
|
document = self.ocr_service.process_image(document)
|
|
self.db.flush()
|
|
return document
|