app/graph/builder.py

"""
app/graph/builder.py
─────────────────────
Graph topology:

 [START] ─► safety ──(blocked)──► output ──► END
               │
           (continue)
               │
             router
              / | \
           rag  |  tool/general
            \   |   /
             ▼  ▼  ▼
              llm
            /       \
    tool_calls?     none
         |            |
   tool_executor      |
         |            |
         ▼            ▼
           memory
              |
            hitl ──(rejected)──► END
              |
          evaluation ──(retry)──► llm
              |
          guardrails
              |
            output
              |
             END
"""

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from app.state import AgentState
from app.nodes.safety import safety_node, safety_route
from app.nodes import (
    router_node, route_selector,
    rag_node,
    llm_node,
    tool_executor_node,
    memory_node,
    hitl_node, hitl_route,
    evaluation_node, eval_route,
    guardrails_node,
    output_node,
)


def build_graph():
    """Compile and return the full LangGraph agent."""
    builder = StateGraph(AgentState)

    # ── Register nodes ────────────────────────────────────────────────────
    builder.add_node("safety",        safety_node)
    builder.add_node("router",        router_node)
    builder.add_node("rag",           rag_node)
    builder.add_node("llm",           llm_node)
    builder.add_node("tool_executor", tool_executor_node)
    builder.add_node("memory",        memory_node)
    builder.add_node("hitl",          hitl_node)
    builder.add_node("evaluation",    evaluation_node)
    builder.add_node("guardrails",    guardrails_node)
    builder.add_node("output",        output_node)

    # ── Entry: safety first ───────────────────────────────────────────────
    builder.set_entry_point("safety")

    # Safety gate — blocked queries skip everything and go straight to output
    builder.add_conditional_edges(
        "safety",
        safety_route,
        {"blocked": "output", "continue": "router"},
    )

    # ── Routing ───────────────────────────────────────────────────────────
    builder.add_conditional_edges(
        "router",
        route_selector,
        {
            "rag":     "rag",   # Knowledge query → retrieve then answer
            "tool":    "llm",   # Tool query → LLM decides which tool to call
            "general": "llm",   # Chat query → straight to LLM
        },
    )

    # RAG retrieval feeds into the LLM node
    builder.add_edge("rag", "llm")

    # After LLM: execute tools if requested, else go straight to memory
    builder.add_conditional_edges(
        "llm",
        lambda s: "tool_executor" if s.get("tool_calls") else "memory",
        {"tool_executor": "tool_executor", "memory": "memory"},
    )

    builder.add_edge("tool_executor", "memory")

    # Memory → HITL review (CHECKPOINT 6)
    builder.add_edge("memory", "hitl")

    # HITL approval gate
    builder.add_conditional_edges(
        "hitl",
        hitl_route,
        {"evaluation": "evaluation", "end": END},
    )

    # Evaluation quality gate — may loop back to LLM (CHECKPOINT 7)
    builder.add_conditional_edges(
        "evaluation",
        eval_route,
        {"retry": "llm", "guardrails": "guardrails"},
    )

    # Safety filter → final output
    builder.add_edge("guardrails", "output")
    builder.add_edge("output", END)

    # MemorySaver persists state across invocations (CHECKPOINT 5)
    checkpointer = MemorySaver()
    return builder.compile(checkpointer=checkpointer)