agents/graph.py

"""
LangGraph Orchestrator Graph — Multi-Agent Suite
=================================================
Assembles all specialist agent nodes into a unified StateGraph with conditional routing edges.

Graph Flow:
    [START] → planner_node
                │
                ├─► (routing == "retrieval") ──► retrieval_node ──┐
                ├─► (routing == "sql") ────────► sql_node ────────┼─► synthesizer_node ──► [END]
                ├─► (routing == "analytics") ──► analytics_node ──┤
                ├─► (routing == "code") ───────► code_node ───────┤
                ├─► (routing == "web") ────────► web_node ────────┤
                │                                                 │
                └─► (routing == "direct") ────────────────────────┘
"""
from __future__ import annotations

import time
from typing import Dict, Any, Literal

import structlog
from langgraph.graph import StateGraph, START, END

from agents.state import CopilotState
from agents.planner import planner_node
from agents.retrieval_agent import retrieval_node
from agents.sql_agent import sql_node
from agents.analytics_agent import analytics_node
from agents.code_agent import code_node
from agents.web_agent import web_node
from agents.synthesizer import synthesizer_node

logger = structlog.get_logger(__name__)


def route_next(state: CopilotState) -> Literal["retrieval", "sql", "analytics", "code", "web", "synthesizer"]:
    """
    Conditional edge function evaluated after supervisor/planner.
    Routes to specialist node based on planner decision.
    """
    decision = state.get("routing_decision", "direct")
    logger.info("Routing conditional edge evaluated", decision=decision)

    if decision == "retrieval":
        return "retrieval"
    elif decision == "sql":
        return "sql"
    elif decision == "analytics":
        return "analytics"
    elif decision == "code":
        return "code"
    elif decision == "web":
        return "web"
    else:
        return "synthesizer"


def build_copilot_graph():
    """
    Constructs and compiles the multi-agent executable LangGraph.
    """
    builder = StateGraph(CopilotState)

    # 1. Add Nodes
    builder.add_node("planner", planner_node)
    builder.add_node("retrieval", retrieval_node)
    builder.add_node("sql", sql_node)
    builder.add_node("analytics", analytics_node)
    builder.add_node("code", code_node)
    builder.add_node("web", web_node)
    builder.add_node("synthesizer", synthesizer_node)

    # 2. Add Edges
    builder.add_edge(START, "planner")

    # Conditional branching from planner
    builder.add_conditional_edges(
        "planner",
        route_next,
        {
            "retrieval": "retrieval",
            "sql": "sql",
            "analytics": "analytics",
            "code": "code",
            "web": "web",
            "synthesizer": "synthesizer",
        },
    )

    # All specialist nodes feed into final synthesizer
    builder.add_edge("retrieval", "synthesizer")
    builder.add_edge("sql", "synthesizer")
    builder.add_edge("analytics", "synthesizer")
    builder.add_edge("code", "synthesizer")
    builder.add_edge("web", "synthesizer")
    builder.add_edge("synthesizer", END)

    # Compile executable graph
    return builder.compile()


# ── Executable Compiled Graph ─────────────────────────────────
copilot_graph = build_copilot_graph()


async def run_copilot_pipeline(
    query: str,
    tenant_id: str,
    user_id: str,
    user_role: str = "employee",
    session_id: str = "default",
    history: list = None,
) -> Dict[str, Any]:
    """
    High-level runner interface for the LangGraph Copilot Pipeline.
    """
    start_time = time.perf_counter()

    initial_state: CopilotState = {
        "query": query,
        "tenant_id": tenant_id,
        "user_id": user_id,
        "user_role": user_role,
        "session_id": session_id,
        "history": history or [],
        "agent_outputs": [],
        "retrieved_chunks": [],
        "citations": [],
    }

    logger.info(
        "Invoking Copilot LangGraph pipeline",
        query=query,
        tenant_id=tenant_id,
        user_id=user_id,
    )

    final_state = await copilot_graph.ainvoke(initial_state)

    duration_ms = (time.perf_counter() - start_time) * 1000

    return {
        "response": final_state.get("final_response", ""),
        "routing_decision": final_state.get("routing_decision", "direct"),
        "agent_type": final_state.get("routing_decision", "synthesizer"),
        "citations": final_state.get("citations", []),
        "sources": final_state.get("citations", []),
        "agent_outputs": final_state.get("agent_outputs", []),
        "tokens_used": final_state.get("tokens_used", 0),
        "latency_ms": round(duration_ms, 2),
        "error": final_state.get("error"),
    }