graph.py

from langgraph.graph import END, START, StateGraph

from src.nodes import (
    final_response_node,
    llm_response_node,
    orchestrator_node,
    route_decision,
    router_node,
    should_call_tool,
    tool_node,
)
from src.state import MAX_ITERS, AgentState


def build_graph():
    builder = StateGraph(AgentState)

    builder.add_node("router",         router_node)
    builder.add_node("llm_response",   llm_response_node)
    builder.add_node("orchestrator",   orchestrator_node)
    builder.add_node("tool_node",      tool_node)
    builder.add_node("final_response", final_response_node)

    builder.add_edge(START, "router")

    builder.add_conditional_edges(
        "router", route_decision,
        {"respond": "llm_response", "orchestrator": "orchestrator"},
    )

    builder.add_edge("llm_response", END)

    builder.add_conditional_edges(
        "orchestrator", should_call_tool,
        {"tool_node": "tool_node", "final_response": "final_response"},
    )

    builder.add_edge("tool_node",      "final_response")
    builder.add_edge("final_response", END)

    return builder.compile()


graph = build_graph()


def run(initial_state: AgentState) -> AgentState:
    """
    Chạy graph với AgentState đã được khởi tạo sẵn.

    Args:
        initial_state: AgentState với conversation_id, sender_id, raw_query đã điền đủ.

    Returns:
        AgentState cuối cùng (có final_answer).
    """
    return graph.invoke(initial_state)


def print_result(result: AgentState, label: str = "") -> None:
    SEP = "═" * 60
    if label:
        print(f"\n{SEP}\n  {label}\n{SEP}")
    print(f"  📌 query_type  : {result.get('query_type')}")
    print(f"  📌 iters       : {result.get('iters')} / {result.get('max_iters')}")
    print(f"  📌 # messages  : {len(result.get('messages', []))}")
    print(f"{'─' * 60}")
    print(f"\n✅ final_answer:\n")
    print(result.get("final_answer") or "(trống)")