app/graph/builder.py

"""
builder.py
──────────
Assembles the LangGraph pipeline: nodes, conditional edges, and the checkpointer.

Flow (matches the original main.py control flow exactly):

    START
      → input_guardrail ──blocked?──► END
                         └─else─────► load_context → classify
    classify ──(conversational|intrinsic)──► agent
             └─(else)──────────────────────► retrieve → agent
    agent → output_guardrail ──blocked?──► END
                              └─else─────► persist → END

Memory: a MongoDB checkpointer keyed by thread_id (= user_id) persists the
clean conversation turns. Falls back to an in-memory saver if Mongo is
unavailable, so the graph always compiles (e.g. in tests / local dev).
"""

from __future__ import annotations

import functools
import inspect
import time

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

from app.config import settings
from app.graph.nodes.agent import agent
from app.graph.nodes.classify import classify
from app.graph.nodes.context import load_context
from app.graph.nodes.guardrails import input_guardrail, output_guardrail
from app.graph.nodes.persist import persist
from app.graph.nodes.retrieve import retrieve
from app.graph.state import GraphState

# ── per-node timing wrapper ──
def _timed(name: str, fn):
    """Wrap a node so each execution logs start + duration. Handles both sync and
    async nodes. Lets you see which node dominates a request's latency."""
    if inspect.iscoroutinefunction(fn):
        @functools.wraps(fn)
        async def awrapper(*args, **kwargs):
            print(f"[timing] ▶ {name} start", flush=True)
            t0 = time.perf_counter()
            try:
                return await fn(*args, **kwargs)
            finally:
                print(f"[timing] ■ {name} done in {time.perf_counter() - t0:.2f}s", flush=True)
        return awrapper

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        print(f"[timing] ▶ {name} start", flush=True)
        t0 = time.perf_counter()
        try:
            return fn(*args, **kwargs)
        finally:
            print(f"[timing] ■ {name} done in {time.perf_counter() - t0:.2f}s", flush=True)
    return wrapper


# Categories that skip RAG (original condition: conversational OR intrinsic).
_SKIP_RAG = ("conversational", "intrinsic")


# ── join barrier ──
def gate(state: GraphState) -> dict:
    """No-op join node: runs only after input_guardrail + load_context + classify
    have all completed (the three parallel pre-steps). Routing happens in
    `route_after_gate`, which reads the merged state from all three."""
    return {}


# ── routing functions ──
def route_after_gate(state: GraphState) -> str:
    """Decide what happens after the parallel pre-steps join.

    • If the input guardrail blocked, stop now — input_guardrail has already set
      the canned refusal response/flags in state.
    • Otherwise route by classification: skip RAG for conversational/intrinsic,
      else go through retrieve first (matches the original control flow)."""
    if state.get("blocked"):
        return "blocked"
    classification = state.get("classification", "")
    if any(tag in classification for tag in _SKIP_RAG):
        return "agent"
    return "retrieve"


def route_after_output_guardrail(state: GraphState) -> str:
    return "blocked" if state.get("blocked") else "continue"


def get_checkpointer():
    """MongoDB checkpointer when configured, else an in-memory fallback."""
    if settings.CONNECTION_STRING and settings.DB_NAME:
        try:
            from langgraph.checkpoint.mongodb import MongoDBSaver

            from app.services.mongo import get_client

            client = get_client()
            if client is not None:
                return MongoDBSaver(
                    client,
                    db_name=settings.DB_NAME,
                    checkpoint_collection_name=settings.CHECKPOINT_COLLECTION,
                )
        except Exception as e:  # noqa: BLE001
            print(f"[checkpoint] MongoDB checkpointer unavailable, using in-memory: {e}")
    return MemorySaver()


def build_graph(checkpointer=None):
    builder = StateGraph(GraphState)

    builder.add_node("input_guardrail", _timed("input_guardrail", input_guardrail))
    builder.add_node("load_context", _timed("load_context", load_context))
    builder.add_node("classify", _timed("classify", classify))
    builder.add_node("gate", _timed("gate", gate))
    builder.add_node("retrieve", _timed("retrieve", retrieve))
    builder.add_node("agent", _timed("agent", agent))
    # builder.add_node("output_guardrail", _timed("output_guardrail", output_guardrail))  # disabled for now (see below)
    builder.add_node("persist", _timed("persist", persist))

    # Fan out: the input guardrail, context load, and classifier are mutually
    # independent (guardrail+classify need only `message`, context needs only
    # `user_id`), so run them concurrently instead of sequentially.
    builder.add_edge(START, "input_guardrail")
    builder.add_edge(START, "load_context")
    builder.add_edge(START, "classify")

    # Join: `gate` runs only after all three parallel steps complete.
    builder.add_edge("input_guardrail", "gate")
    builder.add_edge("load_context", "gate")
    builder.add_edge("classify", "gate")

    # After the join: blocked → END (canned refusal already set by the guardrail),
    # else route by classification.
    builder.add_conditional_edges(
        "gate",
        route_after_gate,
        {"blocked": END, "retrieve": "retrieve", "agent": "agent"},
    )
    builder.add_edge("retrieve", "agent")

    # ── OUTPUT GUARDRAIL: disabled for now ───────────────────────────────────
    # We currently rely on the input guardrail and stream the agent's reply
    # directly (the streaming endpoint can run a "stream-then-guard with retract"
    # check instead — see app/main.py). To put the guardrail back IN the graph,
    # re-add its node above, restore the two edges below, and delete the direct
    # `agent → persist` edge:
    #
    # builder.add_edge("agent", "output_guardrail")
    # builder.add_conditional_edges(
    #     "output_guardrail",
    #     route_after_output_guardrail,
    #     {"blocked": END, "continue": "persist"},
    # )
    builder.add_edge("agent", "persist")
    builder.add_edge("persist", END)

    if checkpointer is None:
        checkpointer = get_checkpointer()
    return builder.compile(checkpointer=checkpointer)


# Process-wide compiled graph (built once, reused per request).
_GRAPH = None


def get_graph():
    global _GRAPH
    if _GRAPH is None:
        _GRAPH = build_graph()
    return _GRAPH