backend/graph.py

# ruff: noqa: E402
from dotenv import load_dotenv

load_dotenv()
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

from agents.supervisor import supervisor_node
from agents.profiler import profiler_node
from agents.researcher import researcher_node
from agents.matcher import matcher_node
from agents.verifier import verifier_node
from agents.timeline import timeline_node
from agents.wizard import wizard_node
from agents.planner import planner_node
from agents.critic import critic_node
from agents.risk_scoring import risk_scoring_node
from agents.document_gap_analyzer import document_gap_analyzer_node
from agents.compliance_guardian import compliance_guardian_node
from schemas import AgentState

# W 2026 roku celujemy w Postgres dla subgrafów, tymczasowo bazowy memory_saver.
memory_saver = MemorySaver()


def create_app():
    workflow = StateGraph(AgentState)

    # Rejestracja wszystkich węzłów (Agentów)
    workflow.add_node("supervisor", supervisor_node)
    workflow.add_node("profiler", profiler_node)
    workflow.add_node("researcher", researcher_node)
    workflow.add_node("matcher", matcher_node)
    workflow.add_node("verifier", verifier_node)
    workflow.add_node("timeline", timeline_node)
    workflow.add_node("wizard", wizard_node)
    workflow.add_node("planner", planner_node)
    workflow.add_node("critic", critic_node)
    workflow.add_node("risk_scoring", risk_scoring_node)
    workflow.add_node("document_gap_analyzer", document_gap_analyzer_node)
    workflow.add_node("compliance_guardian", compliance_guardian_node)

    # Definicja krawędzi wejściowych
    workflow.add_edge(START, "supervisor")

    # Logika warunkowa supervisora jako Global Router
    workflow.add_conditional_edges(
        "supervisor",
        lambda state: state.current_agent,
        {
            "profiler": "profiler",
            "researcher": "researcher",
            "matcher": "matcher",
            "verifier": "verifier",
            "timeline": "timeline",
            "wizard": "wizard",
            "planner": "planner",
            "risk_scoring": "risk_scoring",
            "document_gap_analyzer": "document_gap_analyzer",
            "compliance_guardian": "compliance_guardian",
            "end": END,
        },
    )

    # Conditional routing po Critic (idziemy do wizard jeśli is_approved=False, chyba że przekroczono limit)
    workflow.add_conditional_edges(
        "critic",
        lambda state: "approve"
        if (state.critic_evaluation and state.critic_evaluation.is_approved)
        or state.critic_iterations >= state.max_critic_iterations
        else "wizard",
        {"approve": END, "wizard": "wizard"},
    )

    # W architekturze 2026, Wizard idzie zawsze do Critica (Recenzenta)
    workflow.add_edge("wizard", "critic")

    # Powrotne krawędzie do Supervisora
    workflow.add_edge("profiler", "supervisor")
    workflow.add_edge("researcher", "supervisor")
    workflow.add_edge("matcher", "supervisor")
    workflow.add_edge("verifier", "supervisor")
    workflow.add_edge("timeline", "supervisor")
    workflow.add_edge("planner", "supervisor")
    workflow.add_edge("risk_scoring", "supervisor")
    workflow.add_edge("document_gap_analyzer", "supervisor")
    workflow.add_edge("compliance_guardian", "supervisor")

    # Kompilacja z systemem checkpointów dla Human-in-the-Loop
    return workflow.compile(checkpointer=memory_saver, interrupt_before=["critic"])


app = create_app()