backend/graph.py

"""
LangGraph Orchestrator.
Builds the stateful graph:
  guardian → (blocked?) → guardian_block → END
                       → security → supervisor_router → [crisis | phq9 | therapy | diagnosis] → END

The Guardian agent is the outermost gate — runs before security, before anything.
If it detects abuse / off-topic / vulgar content, the pipeline short-circuits.
"""
from langgraph.graph import StateGraph, END

from backend.state import ClinicalState
from backend.agents.guardian import guardian_node, guardian_block_node
from backend.agents.security import security_node
from backend.agents.crisis import crisis_response_node
from backend.agents.therapy import therapy_node
from backend.agents.diagnosis import diagnosis_node
from backend.agents.evaluation import phq9_node, should_trigger_phq9


# ── Guardian Router — conditional edge after guardian node ──────────
def guardian_router(state: ClinicalState) -> str:
    """Routes to guardian_block (→ END) if blocked, otherwise continues to security."""
    if state.get("guardian_blocked"):
        return "guardian_block"
    return "security"


# ── Supervisor Router — conditional edge after security node ───────
def supervisor_router(state: ClinicalState) -> str:
    if state["crisis_detected"]:
        return "crisis_response"
    if state["mode"] == "therapy" and should_trigger_phq9(state):
        return "phq9_checkin"
    if state["mode"] == "diagnosis":
        return "diagnosis"
    return "therapy"


# ── Build graph ────────────────────────────────────────────────────
def build_graph():
    wf = StateGraph(ClinicalState)

    wf.add_node("guardian",        guardian_node)
    wf.add_node("guardian_block",  guardian_block_node)
    wf.add_node("security",        security_node)
    wf.add_node("crisis_response", crisis_response_node)
    wf.add_node("therapy",         therapy_node)
    wf.add_node("diagnosis",       diagnosis_node)
    wf.add_node("phq9_checkin",    phq9_node)

    # Entry point: Guardian is the first gate
    wf.set_entry_point("guardian")

    # Guardian routing: block or continue
    wf.add_conditional_edges(
        "guardian",
        guardian_router,
        {
            "guardian_block": "guardian_block",
            "security":      "security",
        },
    )

    wf.add_edge("guardian_block", END)

    # Security routing: crisis / phq9 / therapy / diagnosis
    wf.add_conditional_edges(
        "security",
        supervisor_router,
        {
            "crisis_response": "crisis_response",
            "phq9_checkin":    "phq9_checkin",
            "therapy":         "therapy",
            "diagnosis":       "diagnosis",
        },
    )

    wf.add_edge("crisis_response", END)
    wf.add_edge("therapy",         END)
    wf.add_edge("diagnosis",       END)
    wf.add_edge("phq9_checkin",    END)

    return wf.compile()


# Built once at import time — shared across all requests
clinical_graph = build_graph()