Spaces:
Runtime error
Runtime error
Download app/api/routes/batch_upload.py from rahim237cm/jurix-api: direct link, hf CLI and curl.
- Browser
- Download file 14.8 kB
-
https://huggingface.co/spaces/rahim237cm/jurix-api/resolve/main/app/api/routes/batch_upload.py
- Command line
-
hf download hf://spaces/rahim237cm/jurix-api/app/api/routes/batch_upload.py
-
curl -L -o batch_upload.py https://huggingface.co/spaces/rahim237cm/jurix-api/resolve/main/app/api/routes/batch_upload.py
14.8 kB
| """ | |
| Batch upload router with WebSocket support for real-time progress tracking. | |
| Handles: | |
| - Batch PDF upload (multiple files) | |
| - Real-time progress updates via WebSocket | |
| - Status tracking (PENDING → PROCESSING → PUBLISHED/REFUSED) | |
| """ | |
| import asyncio | |
| import logging | |
| import uuid | |
| from datetime import datetime, timezone | |
| from typing import List | |
| from fastapi import ( | |
| APIRouter, | |
| Depends, | |
| File, | |
| HTTPException, | |
| UploadFile, | |
| WebSocket, | |
| WebSocketDisconnect, | |
| ) | |
| from jose import JWTError, jwt | |
| from sqlalchemy import select | |
| from sqlalchemy.ext.asyncio import AsyncSession | |
| from app.core.auth import ALGORITHM, SECRET_KEY, get_current_admin_user | |
| from app.core.config import settings | |
| from app.core.database import AsyncSessionLocal, get_db | |
| from app.models.law import Law | |
| from app.models.user import User | |
| from app.tasks.process_law import process_law_async | |
| logger = logging.getLogger(__name__) | |
| # Références fortes vers les tâches de fond. asyncio ne conserve qu'une référence | |
| # FAIBLE sur les tâches créées par create_task : sans ce registre, une tâche peut | |
| # être collectée en plein traitement et le document rester bloqué en "processing", | |
| # l'exception n'apparaissant qu'au ramasse-miettes. | |
| _background_tasks: set[asyncio.Task] = set() | |
| router = APIRouter(tags=["batch-upload"]) | |
| # WebSocket connection manager | |
| class ConnectionManager: | |
| def __init__(self): | |
| self.active_connections: dict[str, WebSocket] = {} | |
| async def connect(self, session_id: str, websocket: WebSocket): | |
| await websocket.accept() | |
| self.active_connections[session_id] = websocket | |
| def disconnect(self, session_id: str): | |
| if session_id in self.active_connections: | |
| del self.active_connections[session_id] | |
| async def send_progress(self, session_id: str, data: dict): | |
| if session_id in self.active_connections: | |
| try: | |
| await self.active_connections[session_id].send_json(data) | |
| except Exception as e: | |
| # `except:` nu attrapait aussi KeyboardInterrupt et | |
| # SystemExit : un Ctrl-C pendant un envoi etait avale et le | |
| # serveur refusait de s'arreter. La deconnexion reste le bon | |
| # traitement d'un envoi echoue, mais elle se journalise. | |
| logger.debug(f"Envoi WebSocket impossible ({session_id}): {e}") | |
| self.disconnect(session_id) | |
| manager = ConnectionManager() | |
| def _jeton_valide(token: str | None) -> bool: | |
| """ | |
| Verifie un jeton d'acces sans toucher la base. | |
| Volontairement limite a la signature et a l'expiration : ouvrir une session | |
| de base a chaque connexion WebSocket ferait payer une requete a chaque | |
| onglet ouvert, pour un canal qui ne fait que diffuser une progression. Les | |
| ECRITURES, elles, restent protegees par get_current_admin_user sur la route | |
| HTTP correspondante. | |
| """ | |
| if not token: | |
| return False | |
| try: | |
| payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM]) | |
| except JWTError: | |
| return False | |
| return bool(payload.get("sub")) | |
| async def websocket_endpoint( | |
| websocket: WebSocket, session_id: str, token: str | None = None | |
| ): | |
| """ | |
| WebSocket endpoint for real-time upload progress. | |
| Client receives JSON messages: | |
| { | |
| "type": "progress", | |
| "law_id": 123, | |
| "filename": "law.pdf", | |
| "progress": 45, | |
| "status": "processing", | |
| "message": "Extracting text..." | |
| } | |
| """ | |
| assert session_id and isinstance(session_id, str), "session_id must be a non-empty string" | |
| # Le jumeau HTTP de cette route (GET /batch-upload/status) rend 401 sans | |
| # jeton ; ce WebSocket, lui, acceptait n'importe qui et diffusait la | |
| # progression d'imports d'administration. Le navigateur ne peut pas poser | |
| # d'en-tete Authorization sur un WebSocket : le jeton passe donc en | |
| # parametre de requete, et est valide AVANT l'acceptation de la connexion. | |
| if not _jeton_valide(token): | |
| await websocket.close(code=4401, reason="Authentification requise") | |
| return | |
| await manager.connect(session_id, websocket) | |
| MAX_KEEPALIVE_CYCLES = 3600 # ~1 hour at 1s keepalive cadence (NASA Rule 2: bounded loops) | |
| try: | |
| for _ in range(MAX_KEEPALIVE_CYCLES): | |
| # Keep connection alive, bounded to prevent infinite hang | |
| await websocket.receive_text() | |
| # Gracefully close after max cycles | |
| await websocket.close(code=1000, reason="Session timeout (max keepalive reached)") | |
| except WebSocketDisconnect: | |
| pass | |
| finally: | |
| manager.disconnect(session_id) | |
| async def batch_upload( | |
| files: List[UploadFile] = File(...), | |
| session_id: str | None = None, | |
| db: AsyncSession = Depends(get_db), | |
| current_admin: User = Depends(get_current_admin_user), | |
| ): | |
| """ | |
| Batch upload multiple PDF files. | |
| Args: | |
| files: List of PDF files to upload | |
| session_id: WebSocket session ID for progress updates | |
| Returns: | |
| { | |
| "session_id": "uuid", | |
| "total_files": 10, | |
| "created_laws": [{"id": 1, "filename": "law1.pdf"}, ...] | |
| } | |
| """ | |
| assert files is not None and len(files) > 0, "At least one file must be uploaded" | |
| assert all(hasattr(f, 'filename') for f in files), "All uploaded items must be valid files" | |
| # Generate session ID if not provided | |
| if not session_id: | |
| session_id = str(uuid.uuid4()) | |
| # Validate files | |
| for file in files: | |
| if not file.filename or not file.filename.lower().endswith('.pdf'): | |
| raise HTTPException(400, f"File {file.filename or 'unknown'} is not a PDF") | |
| # Check file size | |
| file.file.seek(0, 2) # Seek to end | |
| size = file.file.tell() | |
| file.file.seek(0) # Reset | |
| if size > settings.MAX_UPLOAD_SIZE: | |
| raise HTTPException(400, f"File {file.filename} exceeds {settings.MAX_UPLOAD_SIZE / 1024 / 1024}MB limit") | |
| created_laws = [] | |
| # Process each file | |
| for idx, file in enumerate(files): | |
| try: | |
| # Send progress update | |
| await manager.send_progress(session_id, { | |
| "type": "upload_progress", | |
| "current": idx + 1, | |
| "total": len(files), | |
| "filename": file.filename, | |
| "status": "uploading" | |
| }) | |
| # Read file content | |
| content = await file.read() | |
| # Create law entry with PENDING status | |
| from app.services.file_upload_service import get_upload_service | |
| upload_service = get_upload_service() | |
| file_id = str(uuid.uuid4()) | |
| file_path = upload_service.storage_path / f"{file_id}.pdf" | |
| with open(file_path, "wb") as f: | |
| f.write(content) | |
| law = Law( | |
| reference=f"PENDING-{uuid.uuid4().hex[:8]}", | |
| title=f"PENDING-{uuid.uuid4().hex[:8]}", | |
| type="loi", | |
| # content non vide : le schema LawResponse impose min_length=10, | |
| # et une seule ligne a contenu vide faisait echouer GET /laws | |
| # entier avec une ResponseValidationError 500. | |
| content="Document en attente de traitement.", | |
| status="pending", | |
| # file_id etait genere mais JAMAIS persiste : sans lui, | |
| # /laws/{id}/download, /pdf-info et /page/{n} | |
| # renvoyaient 404 pour tout document issu d'un lot. | |
| file_id=file_id, | |
| original_filename=file.filename, | |
| processing_progress=0, | |
| processing_started_at=None | |
| ) | |
| db.add(law) | |
| await db.flush() # Get ID | |
| created_laws.append({ | |
| "id": law.id, | |
| "filename": file.filename, | |
| "file_id": file_id, | |
| "status": "pending" | |
| }) | |
| # Send creation confirmation | |
| await manager.send_progress(session_id, { | |
| "type": "file_created", | |
| "law_id": law.id, | |
| "filename": file.filename, | |
| "status": "pending" | |
| }) | |
| except Exception as e: | |
| await manager.send_progress(session_id, { | |
| "type": "error", | |
| "filename": file.filename, | |
| "error": str(e) | |
| }) | |
| await db.commit() | |
| # La session `db` vient de Depends(get_db) : elle est fermee des le retour | |
| # du handler. La passer a une tache detachee garantissait un echec. | |
| # process_batch ouvre desormais la sienne. | |
| # | |
| # La reference est conservee dans _background_tasks : asyncio ne garde qu'une | |
| # reference faible sur les taches, une tache non referencee peut donc etre | |
| # collectee en plein vol et le document rester bloque en "processing". | |
| task = asyncio.create_task(process_batch(created_laws, session_id)) | |
| _background_tasks.add(task) | |
| task.add_done_callback(_background_tasks.discard) | |
| return { | |
| "session_id": session_id, | |
| "total_files": len(files), | |
| "created_laws": created_laws | |
| } | |
| async def process_batch(laws: List[dict], session_id: str) -> None: | |
| """ | |
| Traite en arriere-plan tous les fichiers d'un lot. | |
| Ouvre sa PROPRE session par document, et ne recoit plus celle de la requete : | |
| `Depends(get_db)` la ferme des le retour du handler, donc chaque acces | |
| echouait. Une session par document plutot qu'une pour le lot entier : le | |
| traitement dure plusieurs minutes par piece, et garder une connexion du pool | |
| ouverte pendant toute la duree du lot le viderait; par ailleurs un document | |
| en echec n'entraine pas les suivants. | |
| """ | |
| for law_data in laws: | |
| law_id = law_data.get("id") | |
| filename = law_data.get("filename") | |
| try: | |
| async with AsyncSessionLocal() as db: | |
| law = ( | |
| await db.execute(select(Law).where(Law.id == law_id)) | |
| ).scalar_one_or_none() | |
| if not law: | |
| logger.error(f"❌ Loi ID={law_id} introuvable pour le traitement") | |
| continue | |
| law.status = "processing" | |
| law.processing_progress = 0 | |
| # UTC naif : la colonne est un DateTime sans fuseau. | |
| law.processing_started_at = datetime.now(timezone.utc).replace(tzinfo=None) | |
| await db.commit() | |
| await manager.send_progress(session_id, { | |
| "type": "processing_start", | |
| "law_id": law_id, | |
| "filename": filename, | |
| "status": "processing", | |
| }) | |
| result = await process_law_async(law_id, law_data["file_id"]) | |
| succeeded = result.get("status") == "completed" | |
| # Session distincte : le pipeline a pu durer plusieurs minutes. | |
| async with AsyncSessionLocal() as db: | |
| law = ( | |
| await db.execute(select(Law).where(Law.id == law_id)) | |
| ).scalar_one_or_none() | |
| if law: | |
| if succeeded: | |
| law.status = "published" | |
| law.processing_progress = 100 | |
| # processing_error n'est PAS remis a None : le pipeline | |
| # y a ecrit ce qui manque a un document publie (pages | |
| # non extraites, aucun vecteur), le seul signal qu'il | |
| # est incomplet. | |
| else: | |
| law.status = "refused" | |
| law.processing_error = str(result.get("errors", [])) | |
| await db.commit() | |
| await manager.send_progress(session_id, { | |
| "type": "processing_complete", | |
| "law_id": law_id, | |
| "filename": filename, | |
| # etat calcule et non lu sur un objet detache | |
| "status": "published" if succeeded else "refused", | |
| "progress": 100, | |
| }) | |
| except Exception as e: | |
| logger.error(f"❌ Echec du traitement de la loi {law_id}: {e}", exc_info=True) | |
| try: | |
| async with AsyncSessionLocal() as db: | |
| law = ( | |
| await db.execute(select(Law).where(Law.id == law_id)) | |
| ).scalar_one_or_none() | |
| if law: | |
| law.status = "refused" | |
| law.processing_error = str(e) | |
| await db.commit() | |
| except Exception as inner: | |
| # except nu auparavant : il avalait jusqu'aux KeyboardInterrupt | |
| # et masquait totalement l'echec de session ci-dessus. | |
| logger.error(f"❌ Impossible de marquer la loi {law_id} en echec: {inner}") | |
| await manager.send_progress(session_id, { | |
| "type": "processing_error", | |
| "law_id": law_id, | |
| "filename": filename, | |
| "error": str(e), | |
| }) | |
| async def get_batch_status( | |
| status: str | None = None, | |
| db: AsyncSession = Depends(get_db), | |
| current_admin: User = Depends(get_current_admin_user), | |
| ): | |
| """ | |
| Get all laws with optional status filter. | |
| Args: | |
| status: Filter by status (pending, processing, published, refused) | |
| Returns: | |
| List of laws with their processing status | |
| """ | |
| assert status is None or isinstance(status, str), "Status filter must be a string or None" | |
| assert status is None or status in ("pending", "processing", "published", "refused"), \ | |
| f"Invalid status: {status}. Must be one of: pending, processing, published, refused" | |
| query = select(Law) | |
| if status: | |
| query = query.where(Law.status == status) | |
| query = query.order_by(Law.created_at.desc()) | |
| result = await db.execute(query) | |
| laws = result.scalars().all() | |
| return { | |
| "total": len(laws), | |
| "laws": [ | |
| { | |
| "id": law.id, | |
| "title": law.title, | |
| "reference": law.reference, | |
| "status": law.status, | |
| "processing_progress": law.processing_progress, | |
| "processing_error": law.processing_error, | |
| "created_at": law.created_at.isoformat() if law.created_at else None, | |
| "processing_started_at": law.processing_started_at.isoformat() if law.processing_started_at else None | |
| } | |
| for law in laws | |
| ] | |
| } | |