SFCR_Extraction/04 - Scripts/workflow_agents.py

from langgraph.graph import StateGraph, START, END

# ⚠️ Import corrigé (selon version)
try:
    from langgraph.prebuilt import ToolNode
except ImportError:
    # fallback si ToolNode n'existe pas
    from langchain_core.runnables import RunnableLambda

    def ToolNode(tools):
        def run_tools(state):
            return state  # à adapter si besoin
        return RunnableLambda(run_tools)

# Assure-toi que ces imports sont corrects
from Agents import AgentState, agent_extracteur, agent_builder, agent_ocr, tools

workflow = StateGraph(AgentState)

# 1. Définition des Nœuds
workflow.add_node("agent_ocr", agent_ocr)
workflow.add_node("agent_extracteur", agent_extracteur)
workflow.add_node("agent_builder", agent_builder)
workflow.add_node("tools", ToolNode(tools))

# 2. Définition des Arêtes
workflow.add_edge(START, "agent_ocr")
workflow.add_edge("agent_ocr", "agent_extracteur")
workflow.add_edge("agent_extracteur", "agent_builder")


workflow.add_edge("tools",END)


# 5. Compilation
app = workflow.compile() 

# 6. Graph
try:
    with open("graph_workflow.png", "wb") as f:
        f.write(app.get_graph().draw_mermaid_png())
    print(" Graphique du workflow généré sous : graph_workflow.png")
except Exception as e:
    print(f" Erreur génération image : {e}")