SFCR_Extraction/03 - 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]
    page : List[int]
    entreprise_name : str 
    section_name : str
    lignes : 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")
    # On récupère la page unique envoyée par le main
    page_val = state.get("page")

    if page_val is None:
        raise ValueError("page est None dans state")

    # Conversion en entier au cas où
    try:
        page_index = int(page_val) - 1 
    except TypeError:
        # Si c'est une liste [45], on prend le premier élément
        page_index = int(page_val[0]) - 1

    all_points = []
    texte_accumule = []

    try:
        # Traitement de la page unique
        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)

        # Structure spatiale
        points = to_points(data_propre)
        all_points.extend(points)

        # Texte pour le LLM
        lignes_page = formater_donnees_section(data_propre, page_index)
        texte_accumule.extend(lignes_page)

        print(f" Page {page_index + 1} traitée par l'OCR.")

    except Exception as e:
        print(f"Erreur lors de l'OCR Page {page_index + 1}: {e}")
        raise e

    # DBSCAN et Markdown
    clusters = cluster_lines(all_points, eps=0.5)
    lignes = build_lines(clusters)
    lignes = merge_close_lines(lignes)
    markdown = transform_to_clean_markdown(lignes)
    
    contenu_final = "\n".join(texte_accumule)

    return {
        "messages": [HumanMessage(content=f"Voici les données OCR brutes :\n{contenu_final}")],
        "clusters" : clusters , 
        "lignes" : lignes ,
        "markdown": markdown 
    }




def agent_extracteur(state: AgentState):
    prompt = """
    Tu es un extracteur de tableaux QRT (Solvabilité II).
    Objectif : convertir le tableau en JSON sans perte d’information.

    Règles :
    - Utilise Rxxxx (lignes) et Cxxxx (colonnes) comme clés directement.
    - Copie STRICTEMENT les valeurs, sans calcul.
    - Supprime uniquement les espaces dans les nombres : "3 297 388" → 3297388.
    - Garde toute la précision et les 0.
    - Structure attendue : {"Cxxxx": {"Rxxxx": valeur, ...}, ...}

    ⚠️ CONTRAINTE DE SORTIE :
    - Réponds UNIQUEMENT avec l'objet JSON.
    - PAS de markdown (```json), PAS de texte, PAS d'explications.

    Sortie attendue :
    {
        "Cxxxx": { "Rxxxx": 123456, "Ryyyy": 789 }
    }
    """

    markdown_content = state.get("markdown", "")
    if not markdown_content or str(markdown_content).strip() == "":
        backup_content = json.dumps(state.get("lignes", "Aucune donnée trouvée"))
        input_content = f"Note : Le markdown était vide. Voici les lignes brutes :\n{backup_content}"
    else:
        input_content = markdown_content

    msg = [
        SystemMessage(content=prompt),
        HumanMessage(content=json.dumps(input_content))
    ]

    response = model_openai.invoke(msg)
    return {"messages": [response]}





def agent_builder(state: AgentState):
    print(f" Construction du fichier Excel pour : {state['entreprise_name']}...")

    try:
        import json
        import pandas as pd
        import os
        import re

        # 1. Extraction du contenu JSON pur
        content = state["messages"][-1].content
        # Nettoyage au cas où le modèle aurait ajouté des balises markdown malgré les consignes
        content = re.sub(r'```json|```', '', content).strip()
        data_json = json.loads(content)

        # 2. Gestion du chemin de sauvegarde (identique)
        base_outputs = os.path.join("..", "04 - Outputs")
        match = re.search(r'(\d{4})', state['entreprise_name'])
        annee = match.group(1) if match else "2025"
        nom_entreprise = state['entreprise_name'].split('_')[0].replace(" ", "_")
        target_folder = os.path.join(base_outputs, annee, nom_entreprise)
        
        if not os.path.exists(target_folder):
            os.makedirs(target_folder)
        
        output_file = os.path.join(target_folder, f"Rapport_{state['section_name']}.xlsx")

        # 3. Construction du DataFrame à partir de 'data_json' uniquement
        # On récupère l'ensemble unique de tous les Rxxxx présents dans toutes les colonnes
        all_rows = sorted(list(set(r for col in data_json.values() for r in col.keys())))
        all_cols = sorted(list(data_json.keys()))

        df = pd.DataFrame(index=all_rows, columns=all_cols)

        for col, row_values in data_json.items():
            for row, val in row_values.items():
                df.at[row, col] = val

        # 4. Exportation
        with pd.ExcelWriter(output_file, engine="xlsxwriter") as writer:
            df.to_excel(writer, sheet_name="QRT_Export", index=True)
            # Ajoute ici tes formats xlsxwriter si nécessaire

        success_msg = f"✓ Sauvegardé dans : {output_file}"
        print(success_msg)
        return {"messages": [HumanMessage(content=success_msg)]}

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