Định nghĩa các nodes trong graph

nodes.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, Literal
+
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+
+from src.llm import llm
+from src.prompts import (
+    final_response_prompt,
+    orchestrator_prompt,
+    respond_prompt,
+    router_prompt,
+)
+from src.state import MAX_ITERS, AgentState, QueryType
+from src.tools.chat_tools import TOOL_MAP, TOOLS
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
+)
+logger = logging.getLogger("agent_router")
+
+_llm_with_tools = llm.bind_tools(TOOLS)
+
+VALID_LABELS:  set[str]  = {"respond", "summary"}
+DEFAULT_LABEL: QueryType = "respond"
+
+
+# ════════════════════════════════════════════════════════════════════
+# NODE 1 — RouterNode
+# ════════════════════════════════════════════════════════════════════
+def router_node(state: AgentState) -> AgentState:
+    """Validate state, gọi LLM để phân loại query → query_type."""
+    if not state.get("raw_query"):
+        raise ValueError("AgentState thiếu trường 'raw_query'.")
+    if not state.get("sender_id"):
+        raise ValueError("AgentState thiếu trường 'sender_id'.")
+
+    logger.info("[RouterNode] Phân loại query từ %s", state["sender_id"])
+
+    chain     = router_prompt | llm
+    response  = chain.invoke({
+        "sender_id": state["sender_id"],
+        "time":      state.get("time", ""),
+        "raw_query": state["raw_query"],
+    })
+
+    raw_label  = response.content.strip().lower()
+    query_type: QueryType = raw_label if raw_label in VALID_LABELS else DEFAULT_LABEL  # type: ignore
+
+    if raw_label not in VALID_LABELS:
+        logger.warning(
+            "[RouterNode] Nhãn không hợp lệ '%s', fallback → '%s'",
+            raw_label, DEFAULT_LABEL,
+        )
+
+    logger.info("[RouterNode] query_type = %s", query_type)
+    return {**state, "query_type": query_type, "messages": [HumanMessage(content=state["raw_query"])]}
+
+
+def route_decision(state: AgentState) -> Literal["respond", "orchestrator"]:
+    """Conditional edge: router → respond | orchestrator."""
+    return "respond" if state.get("query_type") == "respond" else "orchestrator"
+
+
+# ════════════════════════════════════════════════════════════════════
+# NODE 2 — LLMResponseNode (luồng respond)
+# ════════════════════════════════════════════════════════════════════
+def llm_response_node(state: AgentState) -> AgentState:
+    """Trả lời trực tiếp câu hỏi không cần tra cứu lịch sử."""
+    logger.info("[LLMResponseNode] Trả lời trực tiếp cho %s", state["sender_id"])
+
+    chain    = respond_prompt | llm
+    response = chain.invoke({"sender_id": state["sender_id"], "raw_query": state["raw_query"]})
+    answer   = response.content
+
+    logger.info("[LLMResponseNode] Đã sinh câu trả lời (%d ký tự)", len(answer))
+    return {**state, "messages": [AIMessage(content=answer)], "final_answer": answer}
+
+
+# ════════════════════════════════════════════════════════════════════
+# NODE 3 — OrchestratorNode
+# ════════════════════════════════════════════════════════════════════
+def orchestrator_node(state: AgentState) -> AgentState:
+    """Gọi LLM (đã bind tool) để sinh tool_call. Dừng nếu đạt max_iters."""
+    iters     = state.get("iters", 0)
+    max_iters = state.get("max_iters", MAX_ITERS)
+
+    logger.info("[OrchestratorNode] iters=%d / max_iters=%d", iters, max_iters)
+
+    if iters >= max_iters:
+        logger.info("[OrchestratorNode] Đã đạt max_iters, bỏ qua gọi tool.")
+        return {**state, "iters": iters}
+
+    prompt_value = orchestrator_prompt.format_messages(
+        sender_id=state["sender_id"],
+        time=state.get("time", ""),
+        raw_query=state["raw_query"],
+    )
+
+    response: AIMessage = _llm_with_tools.invoke(prompt_value)  # type: ignore
+    logger.info(
+        "[OrchestratorNode] tool_calls=%s",
+        [tc["name"] for tc in (response.tool_calls or [])],
+    )
+
+    return {**state, "messages": [response], "iters": iters + 1}
+
+
+def should_call_tool(state: AgentState) -> Literal["tool_node", "final_response"]:
+    """Conditional edge: orchestrator → tool_node | final_response."""
+    last_msg = state["messages"][-1] if state["messages"] else None
+    if isinstance(last_msg, AIMessage) and last_msg.tool_calls:
+        return "tool_node"
+    return "final_response"
+
+
+# ════════════════════════════════════════════════════════════════════
+# NODE 4 — ToolNode (custom)
+# ════════════════════════════════════════════════════════════════════
+def _run_tool(tool_name: str, tool_args: dict[str, Any]) -> str:
+    fn = TOOL_MAP.get(tool_name)
+    if fn is None:
+        return f"[Lỗi] Tool '{tool_name}' không tồn tại trong registry."
+    try:
+        result = fn.invoke(tool_args)
+        return result if isinstance(result, str) else json.dumps(result, ensure_ascii=False)
+    except Exception as e:
+        logger.exception("[ToolNode] Tool '%s' gặp lỗi", tool_name)
+        return f"[Lỗi] Tool '{tool_name}': {e}"
+
+
+def tool_node(state: AgentState) -> AgentState:
+    """Thực thi tất cả tool_calls trong AIMessage cuối, sinh ToolMessage(s)."""
+    last_msg = state["messages"][-1]
+
+    if not isinstance(last_msg, AIMessage) or not last_msg.tool_calls:
+        logger.warning("[ToolNode] Không tìm thấy tool_call.")
+        return state
+
+    tool_messages: list[ToolMessage] = []
+    for tc in last_msg.tool_calls:
+        name, args, cid = tc["name"], tc.get("args", {}), tc["id"]
+        logger.info("[ToolNode] Thực thi tool='%s' args=%s", name, args)
+        result = _run_tool(name, args)
+        logger.info("[ToolNode] Tool='%s' → %d ký tự", name, len(result))
+        tool_messages.append(ToolMessage(content=result, tool_call_id=cid, name=name))
+
+    return {**state, "messages": tool_messages}
+
+
+# ════════════════════════════════════════════════════════════════════
+# NODE 5 — FinalResponseNode
+# ════════════════════════════════════════════════════════════════════
+def _extract_tool_results(state: AgentState) -> str:
+    parts = [f"[{m.name}]\n{m.content}" for m in state["messages"] if isinstance(m, ToolMessage)]
+    return "\n\n".join(parts) if parts else "(Không có kết quả từ tool)"
+
+
+def final_response_node(state: AgentState) -> AgentState:
+    """Tổng hợp ToolMessage(s) và sinh câu trả lời cuối cùng."""
+    logger.info("[FinalResponseNode] Tổng hợp câu trả lời...")
+
+    chain    = final_response_prompt | llm
+    response = chain.invoke({
+        "sender_id":    state["sender_id"],
+        "raw_query":    state["raw_query"],
+        "tool_results": _extract_tool_results(state),
+    })
+    answer = response.content
+
+    logger.info("[FinalResponseNode] Hoàn thành (%d ký tự)", len(answer))
+    return {**state, "messages": [AIMessage(content=answer)], "final_answer": answer}