graph.py

from langgraph.graph import END, START, StateGraph

from src.nodes import (
    direct_response_node,
    final_response_node,
    orchestrator_node,
    should_call_tool,
    tool_node,
)
from src.state import AgentState


def build_graph():
    builder = StateGraph(AgentState)

    builder.add_node("orchestrator",    orchestrator_node)
    builder.add_node("tool_node",       tool_node)
    builder.add_node("final_response",  final_response_node)
    builder.add_node("direct_response", direct_response_node)

    builder.add_edge(START, "orchestrator")

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

    builder.add_edge("tool_node",       "final_response")
    builder.add_edge("final_response",  END)
    builder.add_edge("direct_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"  📌 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)")