SFCR_Extraction/04 - Scripts/Agents.py

import os 
from typing import Annotated , Sequence , TypedDict , Optional , List , Any
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langgraph.graph.message import add_messages
from dotenv import load_dotenv  
from langchain_openai import ChatOpenAI 
from langchain_groq import ChatGroq 
import pandas as pd
from langfuse import get_client
import json

from tools import  excel_code_interpreter

from paddleocr import PaddleOCR

from function import (
    preparer_image_zoom_hd, 
    extraire_donnees_ocr, 
    nettoyage_sortie_ocr, 
    formater_donnees_section,
    sauvegarder_fichier_unique,
    to_points , 
    build_lines,
    cluster_lines , 
    merge_close_lines

)

from clean_DBSCAN import (
    transform_to_clean_markdown
)



load_dotenv() 


tools = [excel_code_interpreter]


class AgentState(TypedDict) : 
    messages : Annotated[Sequence[BaseMessage],add_messages]
    pdf_path : str 
    pages : List[int]
    entreprise_name : str 
    section_name : str
    markdown : str


model_llama = ChatGroq(model="llama-3.3-70b-versatile") 

model_openai = ChatOpenAI(model="gpt-4.1" , temperature=0.2) 

model_ai = model_llama.bind_tools(tools)

ocr = PaddleOCR(use_angle_cls=True, lang='fr', det_limit_side_len=10000, show_log=False)



def agent_ocr(state: AgentState):
    pdf_path = state.get("pdf_path")
    pages = state.get("pages")

    

    texte_accumule = []
    all_points = []

    for page_index in pages:

        page_index = page_index -1
        
        try:
            img_finale = preparer_image_zoom_hd(pdf_path, page_index)
            raw_data = extraire_donnees_ocr(img_finale, ocr)

            data_propre = nettoyage_sortie_ocr(raw_data)

            # 🔥 garder structure brute
            points = to_points(data_propre)
            all_points.extend(points)

            # debug texte
            lignes_page = formater_donnees_section(data_propre, page_index)
            texte_accumule.extend(lignes_page)

            print(f"✅ Page {page_index + 1} traitée.")

        except Exception as e:
            print(f"❌ Erreur Page {page_index + 1}: {e}")

    # 🔥 DBSCAN GLOBAL (beaucoup plus robuste)
    clusters = cluster_lines(all_points, eps=0.5)
    lignes = build_lines(clusters)
    lignes = merge_close_lines(lignes)

    # 🔥 TEXTE pour LLM
    contenu_final = "\n".join(texte_accumule)

    message_ocr = HumanMessage(
        content=f"Voici les données OCR brutes :\n{contenu_final}"
    )
  
    markdown = transform_to_clean_markdown(lignes)
    
    

    return {
        "messages": [message_ocr],
        "lignes" : lignes,
        "markdown": markdown 
    } 




def agent_extracteur(state:AgentState) : 

    prompt = """
    Tu es un extracteur de tableaux QRT (Solvabilité II).

    Ton but est de transcrire les coordonnées spatiales OCR en une hiérarchie logique de données.

    Objectif : convertir le tableau en JSON sans perte d’information.

    Règles :
    - Utilise Rxxxx (lignes) et Cxxxx (colonnes) comme clés.
    - Copie STRICTEMENT les valeurs, sans calcul ni modification.
    - Supprime uniquement les espaces dans les nombres :
    "3 297 388" → 3297388
    - Ne tronque pas, n’arrondis pas, garde toute la précision.
    - Garde les 0.
    - Si pas de libellé pour un Rxxxx → "".
    - Optimise la structure (par lignes ou colonnes) pour réduire la taille.

    ⚠️ CONTRAINTE DE SORTIE (OBLIGATOIRE) :
    - Réponds UNIQUEMENT avec un objet JSON valide
    - NE PAS ajouter de texte avant ou après
    - NE PAS utiliser ```json ou ```
    - NE PAS commenter
    - NE PAS expliquer

    Sortie attendue (exemple strict) :
    {
    "labels": { "Rxxxx": "Rxxxx" },
    "data": {
        "Cxxxx": { "Rxxxx": 123456 }
    }
    }
    """

    msg = [
    SystemMessage(content=prompt),
    HumanMessage(content=json.dumps(state["markdown"]))
     ]

    response = model_llama.invoke(msg)

    return{"messages" : [response]}





def agent_builder(state: AgentState):
    """
    Agent Builder (Version non-IA) : 
    Prend le JSON de l'extracteur et génère physiquement le fichier Excel.
    """
    print(f"🏗️ Construction du fichier Excel pour : {state['entreprise_name']}...")

    try:
        # 1. Récupération du JSON depuis le dernier message de l'extracteur
        # On suppose que l'extracteur a renvoyé un dictionnaire ou une string JSON
        import json
        last_message_content = state["messages"][-1].content
        
        # Nettoyage si le LLM a mis des balises ```json
        if "```json" in last_message_content:
            last_message_content = last_message_content.split("```json")[1].split("```")[0].strip()
        
        payload = json.loads(last_message_content)
        
        labels = payload.get("labels", {})
        data_json = payload.get("data", {})

        # 2. Préparation de la structure du dossier
        folder = state['entreprise_name'].replace(" ", "_")
        if not os.path.exists(folder):
            os.makedirs(folder)
        
        output_file = os.path.join(folder, f"Rapport_{state['section_name']}.xlsx")

        # 3. Construction du DataFrame
        # On récupère toutes les colonnes et toutes les lignes uniques
        all_rows = list(labels.keys())
        all_cols = list(data_json.keys())

        # Création d'une matrice vide remplie de NaN
        df = pd.DataFrame(index=all_rows, columns=all_cols)

        # Remplissage des données
        for col, row_values in data_json.items():
            for row, val in row_values.items():
                df.at[row, col] = val

        # 4. Ajout de la colonne Libellé au début
        #df.insert(0, "Libellé", df.index.map(labels))

        # 5. Exportation avec formatage pro via XlsxWriter
        with pd.ExcelWriter(output_file, engine="xlsxwriter") as writer:
            df.to_excel(writer, sheet_name="QRT_Export", index=True)
            
            workbook  = writer.book
            worksheet = writer.sheets["QRT_Export"]

            # --- Ajout de formats ---
            header_fmt = workbook.add_format({'bold': True, 'bg_color': '#1F3864', 'font_color': 'white', 'border': 1})
            num_fmt    = workbook.add_format({'num_format': '#,##0', 'border': 1})
            label_fmt  = workbook.add_format({'bg_color': '#D6E4F0', 'border': 1})

            # Appliquer les formats aux colonnes
            worksheet.set_column(1, 1, 60) # Libellé
            worksheet.set_column(0, 0, 15) # Code R
            
            worksheet.set_column(2, len(df.columns)+1, 18, num_fmt) # Valeurs

        success_msg = f"✓ Fichier Excel généré : {output_file} ({len(df)} lignes x {len(all_cols)} colonnes)"
        return {"messages": [HumanMessage(content=success_msg)]}

    except Exception as e:
        error_msg = f"❌ Erreur lors de la construction Excel : {str(e)}"
        print(error_msg)
        return {"messages": [HumanMessage(content=error_msg)]}