4NK_IA_back/services/worker/worker.py

"""
Worker Celery pour le pipeline de traitement des documents notariaux
"""
import os
import time
import logging
from celery import Celery
from celery.signals import task_prerun, task_postrun, task_failure
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from pipelines import preprocess, ocr, classify, extract, index, checks, finalize
from utils.database import Document, ProcessingLog, init_db
from utils.storage import get_document, store_artifact

# Configuration du logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration Celery
app = Celery(
    'worker',
    broker=os.getenv("REDIS_URL", "redis://localhost:6379/0"),
    backend=os.getenv("REDIS_URL", "redis://localhost:6379/0")
)

# Configuration de la base de données
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+psycopg://notariat:notariat_pwd@localhost:5432/notariat")
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

@app.task(bind=True, name='pipeline.run')
def pipeline_run(self, doc_id: str):
    """
    Pipeline principal de traitement d'un document
    """
    db = SessionLocal()
    ctx = {"doc_id": doc_id, "db": db}

    try:
        logger.info(f"Début du traitement du document {doc_id}")

        # Mise à jour du statut
        document = db.query(Document).filter(Document.id == doc_id).first()
        if not document:
            raise ValueError(f"Document {doc_id} non trouvé")

        document.status = "processing"
        db.commit()

        # Exécution des étapes du pipeline
        steps = [
            ("preprocess", preprocess.run),
            ("ocr", ocr.run),
            ("classify", classify.run),
            ("extract", extract.run),
            ("index", index.run),
            ("checks", checks.run),
            ("finalize", finalize.run)
        ]

        for step_name, step_func in steps:
            try:
                logger.info(f"Exécution de l'étape {step_name} pour le document {doc_id}")

                # Enregistrement du début de l'étape
                log_entry = ProcessingLog(
                    document_id=doc_id,
                    step_name=step_name,
                    status="started"
                )
                db.add(log_entry)
                db.commit()

                start_time = time.time()

                # Exécution de l'étape
                step_func(doc_id, ctx)

                # Enregistrement de la fin de l'étape
                duration = int((time.time() - start_time) * 1000)  # en millisecondes
                log_entry.status = "completed"
                log_entry.completed_at = time.time()
                log_entry.duration = duration
                db.commit()

                logger.info(f"Étape {step_name} terminée pour le document {doc_id} en {duration}ms")

            except Exception as e:
                logger.error(f"Erreur dans l'étape {step_name} pour le document {doc_id}: {e}")

                # Enregistrement de l'erreur
                log_entry.status = "failed"
                log_entry.completed_at = time.time()
                log_entry.error_message = str(e)
                db.commit()

                # Ajout de l'erreur au document
                if not document.errors:
                    document.errors = []
                document.errors.append(f"{step_name}: {str(e)}")
                document.status = "failed"
                db.commit()

                raise

        # Succès complet
        document.status = "completed"
        db.commit()

        logger.info(f"Traitement terminé avec succès pour le document {doc_id}")

        return {
            "doc_id": doc_id,
            "status": "completed",
            "processing_steps": ctx.get("processing_steps", {}),
            "extracted_data": ctx.get("extracted_data", {})
        }

    except Exception as e:
        logger.error(f"Erreur fatale lors du traitement du document {doc_id}: {e}")

        # Mise à jour du statut d'erreur
        document = db.query(Document).filter(Document.id == doc_id).first()
        if document:
            document.status = "failed"
            if not document.errors:
                document.errors = []
            document.errors.append(f"Erreur fatale: {str(e)}")
            db.commit()

        raise
    finally:
        db.close()

@app.task(name='queue.process_imports')
def process_import_queue():
    """
    Traitement de la queue d'import Redis
    """
    import redis
    import json

    r = redis.Redis.from_url(os.getenv("REDIS_URL", "redis://localhost:6379/0"))

    try:
        # Récupération d'un élément de la queue
        result = r.brpop("queue:import", timeout=1)

        if result:
            _, payload_str = result
            payload = json.loads(payload_str)
            doc_id = payload["doc_id"]

            logger.info(f"Traitement du document {doc_id} depuis la queue")

            # Lancement du pipeline
            pipeline_run.delay(doc_id)

            # Décrémentation du compteur
            r.decr("stats:pending_tasks")

    except Exception as e:
        logger.error(f"Erreur lors du traitement de la queue d'import: {e}")

# Configuration des signaux Celery
@task_prerun.connect
def task_prerun_handler(sender=None, task_id=None, task=None, args=None, kwargs=None, **kwds):
    """Handler avant exécution d'une tâche"""
    logger.info(f"Début de la tâche {task.name} (ID: {task_id})")

@task_postrun.connect
def task_postrun_handler(sender=None, task_id=None, task=None, args=None, kwargs=None, retval=None, state=None, **kwds):
    """Handler après exécution d'une tâche"""
    logger.info(f"Fin de la tâche {task.name} (ID: {task_id}) - État: {state}")

@task_failure.connect
def task_failure_handler(sender=None, task_id=None, exception=None, traceback=None, einfo=None, **kwds):
    """Handler en cas d'échec d'une tâche"""
    logger.error(f"Échec de la tâche {sender.name} (ID: {task_id}): {exception}")

if __name__ == '__main__':
    # Initialisation de la base de données
    init_db()

    # Démarrage du worker
    app.start()