4NK_IA_back/services/worker/tasks/extraction_tasks.py

"""
Tâches d'extraction d'entités des documents
"""
import logging
from typing import Dict, Any, List
from services.worker.celery_app import app

logger = logging.getLogger(__name__)

@app.task(bind=True, name='extraction.extract_entities')
def extract_entities(self, doc_id: str, text: str, doc_type: str, context: Dict[str, Any]) -> Dict[str, Any]:
    """
    Extraction d'entités d'un document

    Args:
        doc_id: Identifiant du document
        text: Texte extrait du document
        doc_type: Type de document classifié
        context: Contexte de traitement

    Returns:
        Résultat de l'extraction d'entités
    """
    try:
        logger.info(f"Début de l'extraction d'entités pour le document {doc_id}")

        # Mise à jour du statut
        self.update_state(
            state='PROGRESS',
            meta={'current_step': 'entity_extraction', 'progress': 0}
        )

        # TODO: Implémenter l'extraction réelle avec LLM
        # Pour l'instant, simulation
        import time
        time.sleep(2)  # Simulation du traitement

        # Extraction simulée basée sur le type de document
        entities = {}

        if doc_type == 'acte_vente':
            entities = {
                'vendeur': {
                    'nom': 'Dupont',
                    'prenom': 'Jean',
                    'adresse': '123 Rue de la Paix, 75001 Paris'
                },
                'acheteur': {
                    'nom': 'Martin',
                    'prenom': 'Marie',
                    'adresse': '456 Avenue des Champs, 75008 Paris'
                },
                'bien': {
                    'adresse': '789 Boulevard Saint-Germain, 75006 Paris',
                    'surface': '85 m²',
                    'prix': '450000 €'
                },
                'notaire': {
                    'nom': 'Durand',
                    'etude': 'Etude Durand & Associés'
                }
            }
        elif doc_type == 'cni':
            entities = {
                'personne': {
                    'nom': 'Dupont',
                    'prenom': 'Jean',
                    'date_naissance': '1985-03-15',
                    'lieu_naissance': 'Paris',
                    'nationalite': 'Française'
                },
                'document': {
                    'numero': '123456789',
                    'pays': 'France',
                    'date_emission': '2020-01-15',
                    'date_expiration': '2030-01-15'
                }
            }
        else:
            entities = {
                'personnes': [],
                'adresses': [],
                'montants': [],
                'dates': []
            }

        result = {
            'doc_id': doc_id,
            'status': 'completed',
            'entities': entities,
            'confidence': 0.85,
            'extraction_method': 'llm_simulation',
            'processing_time': 2.0
        }

        logger.info(f"Extraction d'entités terminée pour le document {doc_id}")
        return result

    except Exception as e:
        logger.error(f"Erreur lors de l'extraction d'entités du document {doc_id}: {e}")
        raise

@app.task(name='extraction.batch_extract')
def batch_extract_entities(doc_ids: list, texts: list, doc_types: list) -> Dict[str, Any]:
    """
    Extraction d'entités en lot

    Args:
        doc_ids: Liste des identifiants de documents
        texts: Liste des textes correspondants
        doc_types: Liste des types de documents correspondants

    Returns:
        Résultats de l'extraction en lot
    """
    if len(doc_ids) != len(texts) or len(doc_ids) != len(doc_types):
        raise ValueError("Le nombre de documents, textes et types doit être identique")

    logger.info(f"Extraction d'entités en lot de {len(doc_ids)} documents")

    results = []
    for doc_id, text, doc_type in zip(doc_ids, texts, doc_types):
        try:
            result = extract_entities.delay(doc_id, text, doc_type, {}).get()
            results.append(result)
        except Exception as e:
            logger.error(f"Erreur lors de l'extraction en lot pour {doc_id}: {e}")
            results.append({
                'doc_id': doc_id,
                'status': 'failed',
                'error': str(e)
            })

    return {
        'batch_status': 'completed',
        'total_documents': len(doc_ids),
        'results': results
    }