308 lines
9.5 KiB
Python
308 lines
9.5 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import Annotated
|
|
|
|
from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile, status
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.config import settings
|
|
from app.core.database import get_db
|
|
from app.core.dependencies import CurrentUser
|
|
from app.core.exceptions import FileSizeError, UnsupportedFileTypeError
|
|
from app.core.logging_config import get_logger
|
|
from app.models.document import Document
|
|
from app.repositories.document_repository import DocumentRepository
|
|
from app.schemas.common import PaginatedResponse, SuccessResponse
|
|
from app.schemas.document import (
|
|
DocumentListResponse,
|
|
DocumentResponse,
|
|
DocumentUploadResponse,
|
|
TemplateMatchRequest,
|
|
TemplateMatchResponse,
|
|
)
|
|
from app.storage.provider import LocalStorageProvider, get_storage_provider
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
router = APIRouter(prefix="/documents", tags=["Documents"])
|
|
|
|
ALLOWED_CONTENT_TYPES = {
|
|
"image/jpeg": "jpg",
|
|
"image/png": "png",
|
|
"image/tiff": "tiff",
|
|
"application/pdf": "pdf",
|
|
}
|
|
|
|
ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".tiff", ".tif", ".pdf"}
|
|
|
|
|
|
def _validate_file(file: UploadFile) -> str:
|
|
"""Validate uploaded file type and size. Returns the content type."""
|
|
if not file.filename:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="Filename is required",
|
|
)
|
|
|
|
# Check extension
|
|
from pathlib import Path
|
|
ext = Path(file.filename).suffix.lower()
|
|
if ext not in ALLOWED_EXTENSIONS:
|
|
raise UnsupportedFileTypeError(ext)
|
|
|
|
# Determine content type
|
|
content_type = file.content_type or ""
|
|
if content_type not in ALLOWED_CONTENT_TYPES:
|
|
# Try to infer from extension
|
|
ext_to_ct = {
|
|
".jpg": "image/jpeg",
|
|
".jpeg": "image/jpeg",
|
|
".png": "image/png",
|
|
".tiff": "image/tiff",
|
|
".tif": "image/tiff",
|
|
".pdf": "application/pdf",
|
|
}
|
|
content_type = ext_to_ct.get(ext, "")
|
|
if not content_type:
|
|
raise UnsupportedFileTypeError(file.content_type or "unknown")
|
|
|
|
return content_type
|
|
|
|
|
|
@router.post(
|
|
"/upload",
|
|
response_model=DocumentUploadResponse,
|
|
status_code=status.HTTP_201_CREATED,
|
|
summary="Upload Document",
|
|
description="Upload a document (JPG, JPEG, PNG, TIFF, or PDF) for processing.",
|
|
)
|
|
async def upload_document(
|
|
file: UploadFile = File(..., description="Document file to upload"),
|
|
current_user: CurrentUser = None,
|
|
db: Session = Depends(get_db),
|
|
) -> DocumentUploadResponse:
|
|
"""Upload a document for processing."""
|
|
content_type = _validate_file(file)
|
|
|
|
# Read file data
|
|
file_data = await file.read()
|
|
|
|
# Check file size
|
|
if len(file_data) > settings.storage_max_file_size_bytes:
|
|
raise FileSizeError(settings.storage_max_file_size_mb)
|
|
|
|
# Store file
|
|
storage = get_storage_provider()
|
|
checksum = storage.compute_checksum(file_data)
|
|
safe_filename = file.filename or "unknown"
|
|
stored_filename = storage.generate_filename(safe_filename)
|
|
storage_path = storage.save_file(file_data, "documents", stored_filename)
|
|
|
|
# Create document record
|
|
doc_repo = DocumentRepository(db)
|
|
document = Document(
|
|
filename=stored_filename,
|
|
original_filename=safe_filename,
|
|
content_type=content_type,
|
|
file_size=len(file_data),
|
|
checksum=checksum,
|
|
storage_path=storage_path,
|
|
status="pending",
|
|
uploaded_by=current_user.id if current_user else None,
|
|
)
|
|
doc_repo.create(document)
|
|
db.commit()
|
|
db.refresh(document)
|
|
|
|
logger.info(
|
|
"document_uploaded",
|
|
document_id=str(document.id),
|
|
filename=safe_filename,
|
|
size=len(file_data),
|
|
content_type=content_type,
|
|
)
|
|
|
|
# Trigger async processing via Celery
|
|
try:
|
|
from app.tasks.document_tasks import process_document_task
|
|
process_document_task.delay(str(document.id))
|
|
except Exception as e:
|
|
logger.warning("celery_dispatch_failed", error=str(e), document_id=str(document.id))
|
|
|
|
return DocumentUploadResponse.model_validate(document)
|
|
|
|
|
|
@router.get(
|
|
"/{document_id}",
|
|
response_model=DocumentResponse,
|
|
summary="Get Document",
|
|
description="Retrieve a document by its ID with all extracted content.",
|
|
)
|
|
def get_document(
|
|
document_id: uuid.UUID,
|
|
current_user: CurrentUser = None,
|
|
db: Session = Depends(get_db),
|
|
) -> DocumentResponse:
|
|
"""Get a document by ID."""
|
|
doc_repo = DocumentRepository(db)
|
|
document = doc_repo.get_with_pages(document_id)
|
|
if not document:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail=f"Document '{document_id}' not found",
|
|
)
|
|
return DocumentResponse.model_validate(document)
|
|
|
|
|
|
@router.get(
|
|
"",
|
|
response_model=PaginatedResponse[DocumentListResponse],
|
|
summary="List Documents",
|
|
description="List all documents with pagination.",
|
|
)
|
|
def list_documents(
|
|
page: int = Query(default=1, ge=1),
|
|
page_size: int = Query(default=20, ge=1, le=100),
|
|
status_filter: str | None = Query(default=None, alias="status"),
|
|
current_user: CurrentUser = None,
|
|
db: Session = Depends(get_db),
|
|
) -> PaginatedResponse[DocumentListResponse]:
|
|
"""List documents with pagination and optional status filter."""
|
|
doc_repo = DocumentRepository(db)
|
|
offset = (page - 1) * page_size
|
|
filters = {}
|
|
if status_filter:
|
|
filters["status"] = status_filter
|
|
|
|
documents = doc_repo.get_all(
|
|
offset=offset,
|
|
limit=page_size,
|
|
filters=filters,
|
|
order_by="created_at",
|
|
order_desc=True,
|
|
)
|
|
total = doc_repo.count(filters=filters)
|
|
|
|
items = [DocumentListResponse.model_validate(doc) for doc in documents]
|
|
return PaginatedResponse.create(
|
|
items=items,
|
|
total=total,
|
|
page=page,
|
|
page_size=page_size,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/{document_id}/template",
|
|
response_model=list[TemplateMatchResponse],
|
|
summary="Get Document Template Matches",
|
|
description="Get template matching results for a document.",
|
|
)
|
|
def get_document_template_matches(
|
|
document_id: uuid.UUID,
|
|
current_user: CurrentUser = None,
|
|
db: Session = Depends(get_db),
|
|
) -> list[TemplateMatchResponse]:
|
|
"""Get template matches for a document."""
|
|
from app.repositories.document_repository import TemplateMatchRepository
|
|
|
|
doc_repo = DocumentRepository(db)
|
|
document = doc_repo.get_by_id(document_id)
|
|
if not document:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail=f"Document '{document_id}' not found",
|
|
)
|
|
|
|
match_repo = TemplateMatchRepository(db)
|
|
matches = match_repo.get_document_matches(document_id)
|
|
|
|
results = []
|
|
for match in matches:
|
|
resp = TemplateMatchResponse(
|
|
id=match.id,
|
|
document_id=match.document_id,
|
|
format_id=match.format_id,
|
|
confidence_score=match.confidence_score,
|
|
match_details=match.match_details,
|
|
selected=match.selected,
|
|
template_name=match.template.name if match.template else None,
|
|
created_at=match.created_at,
|
|
)
|
|
results.append(resp)
|
|
|
|
return results
|
|
|
|
|
|
@router.post(
|
|
"/{document_id}/extraction",
|
|
response_model=dict,
|
|
summary="Extract Document Data",
|
|
description="Match a document against templates and extract key-value & table data.",
|
|
)
|
|
def extract_document_data(
|
|
document_id: uuid.UUID,
|
|
current_user: CurrentUser = None,
|
|
db: Session = Depends(get_db),
|
|
) -> dict:
|
|
"""Extract document data using matched template mappings."""
|
|
doc_repo = DocumentRepository(db)
|
|
document = doc_repo.get_by_id(document_id)
|
|
if not document:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail=f"Document '{document_id}' not found",
|
|
)
|
|
|
|
if document.status != "completed":
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail=f"Document must be in 'completed' status. Current status: '{document.status}'",
|
|
)
|
|
|
|
from app.services.extraction_service import ExtractionService
|
|
extraction_service = ExtractionService(db)
|
|
try:
|
|
result = extraction_service.extract_document_data(document_id)
|
|
return result
|
|
except Exception as e:
|
|
logger.exception("extraction_endpoint_failed", document_id=str(document_id), error=str(e))
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail=f"Extraction failed: {str(e)}",
|
|
)
|
|
|
|
|
|
@router.delete(
|
|
"/{document_id}",
|
|
response_model=SuccessResponse,
|
|
summary="Delete Document",
|
|
description="Delete a document and its associated data.",
|
|
)
|
|
def delete_document(
|
|
document_id: uuid.UUID,
|
|
current_user: CurrentUser = None,
|
|
db: Session = Depends(get_db),
|
|
) -> SuccessResponse:
|
|
"""Delete a document."""
|
|
doc_repo = DocumentRepository(db)
|
|
document = doc_repo.get_by_id(document_id)
|
|
if not document:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail=f"Document '{document_id}' not found",
|
|
)
|
|
|
|
# Delete stored file
|
|
try:
|
|
storage = get_storage_provider()
|
|
storage.delete_file(document.storage_path)
|
|
except Exception as e:
|
|
logger.warning("file_delete_failed", error=str(e), path=document.storage_path)
|
|
|
|
doc_repo.delete(document)
|
|
db.commit()
|
|
|
|
return SuccessResponse(message=f"Document '{document_id}' deleted successfully")
|