graph.py

from datetime import datetime, timezone, timedelta
from typing import Literal

from pydantic import BaseModel
from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage, AIMessage
from langgraph.graph import StateGraph, END

from src.state import State
from src.llm import reasoning_llm, non_reasoning_llm
from src.prompts import (
    ROUTER_SYSTEM_PROMPT,
    ROUTER_MEMORY_CONTEXT,
    ORCHESTRATOR_SYSTEM_PROMPT,
    RESPONSE_SYSTEM_PROMPT,
)
from src.tools import ORCHESTRATOR_TOOLS, TOOL_REGISTRY
from src.memory_store import get_memory, get_prev_docs
from src.source_display import display_source_from_metadata

MAX_ITERS = 2
_UTC7 = timezone(timedelta(hours=7))


def _get_current_time() -> str:
    return datetime.now(_UTC7).strftime("%d/%m/%Y %H:%M:%S (UTC+7)")


class RouterDecision(BaseModel):
    needs_data: bool
    reuse_prev_docs: bool
    response_type: Literal["rag", "greeting", "out_of_scope"]


_router_llm = non_reasoning_llm.with_structured_output(RouterDecision)
_orchestrator_llm = reasoning_llm.bind_tools(ORCHESTRATOR_TOOLS)


def router_node(state: State) -> dict:
    # LẤY USDER ID
    user_id = state["user_id"]
    session_id = state.get("session_id")

    # LẤY CUỘC NÓI CŨ CỦA USER ID ĐÓ
    memory_vars = get_memory(user_id, session_id=session_id).load_memory_variables({})
    conversation_history = memory_vars.get("history") or "(chưa có lịch sử)"

    # DỰA VÀO USER ID ĐỂ TÌM KIẾM NHỮNG DOCS CHỨA CONTEXT CỦA DỮ LIỆU CŨ
    prev_docs = get_prev_docs(user_id, session_id=session_id)
    if prev_docs:
        doc_titles = "\n".join(
            f"  - {d.metadata.get('title') or d.metadata.get('source', 'N/A')}"
            for d in prev_docs
        )
    else:
        doc_titles = "  (không có)"

    # TỔNG HỢP LẠI DOCS VÀ CUỘC NÓI CHUYỆN CŨ THÀNH CONTEXT
    memory_context = ROUTER_MEMORY_CONTEXT.format(
        conversation_history=conversation_history,
        doc_count=len(prev_docs),
        doc_titles=doc_titles,
    )

    decision: RouterDecision = _router_llm.invoke([
        SystemMessage(content=ROUTER_SYSTEM_PROMPT.format(
            current_time=_get_current_time(),
            memory_context=memory_context,
        )),
        HumanMessage(content=state["query"]),
    ])

    print(f"[ROUTER] response_type={decision.response_type} | needs_data={decision.needs_data} | reuse_prev_docs={decision.reuse_prev_docs}")

    # Determine initial raw_docs for this turn
    if decision.reuse_prev_docs and prev_docs:
        initial_docs = prev_docs
    else:
        initial_docs = []

    return {
        "needs_data": decision.needs_data,
        "response_type": decision.response_type,
        "rewritten_query": state["query"],
        "messages": [HumanMessage(content=state["query"])],
        "raw_docs": initial_docs,
    }


def rule_response_node(state: State) -> dict:
    response_type = state.get("response_type") or "out_of_scope"
    if response_type == "greeting":
        content = "Chào bạn, mình có thể hỗ trợ gì cho bạn ?"
    else:
        content = (
            "Mình chỉ là một trợ lý ảo hỗ trợ các khúc mắc liên quan đến các thông tin của tập đoàn."
        )
    return {"messages": [AIMessage(content=content)]}


def orchestrator_node(state: State) -> dict:
    system = ORCHESTRATOR_SYSTEM_PROMPT.format(current_time=_get_current_time())
    messages = [SystemMessage(content=system)] + state["messages"]
    response = _orchestrator_llm.invoke(messages)
    tool_calls = getattr(response, "tool_calls", [])
    print(f"[ORCHESTRATOR] tool_calls={[tc['name'] for tc in tool_calls]}")
    return {"messages": [response]}


def tool_node(state: State) -> dict:
    last_message = state["messages"][-1]
    existing_docs = state.get("raw_docs") or []
    tool_messages = []
    new_docs = []
    rewritten_query = state.get("rewritten_query") or state["query"]

    for tool_call in last_message.tool_calls:
        fn = TOOL_REGISTRY.get(tool_call["name"])
        if fn is None:
            print(f"[TOOL] Tool not found: {tool_call['name']}")
            continue
        args = dict(tool_call["args"])
        if tool_call["name"] == "vector_search":
            args["search_method"] = state.get("search_method") or "hybrid"
            if args.get("query"):
                rewritten_query = args["query"]
        print(f"[TOOL] Calling {tool_call['name']} with args={args}")
        docs = fn(**args)
        print(f"[TOOL] {tool_call['name']} returned {len(docs)} docs")
        new_docs.extend(docs)
        tool_messages.append(
            ToolMessage(
                content=f"Đã tìm thấy {len(docs)} tài liệu.",
                tool_call_id=tool_call["id"],
            )
        )

    return {
        "messages": tool_messages,
        "raw_docs": existing_docs + new_docs,
        "iter_count": (state.get("iter_count") or 0) + 1,
        "rewritten_query": rewritten_query,
    }


def _format_docs(docs: list) -> str:
    if not docs:
        return ""
    parts = []
    for i, doc in enumerate(docs, 1):
        source = display_source_from_metadata(doc.metadata)
        title = doc.metadata.get("title", "")
        header = f"Tài liệu {i}" + (f" — {title}" if title else "")
        body = f"{header}:\n{doc.page_content}"
        if source:
            body += f"\nNguồn: {source}"
        parts.append(body)
    return "\n\n".join(parts)


def response_node(state: State) -> dict:
    raw_docs = state.get("raw_docs") or []
    print(f"[RESPONSE] raw_docs count={len(raw_docs)}")
    formatted_docs = _format_docs(raw_docs)
    query = state.get("rewritten_query") or state["query"]
    user_id = state["user_id"]
    session_id = state.get("session_id")
    memory_vars = get_memory(user_id, session_id=session_id).load_memory_variables({})
    conversation_history = memory_vars.get("history") or "(chưa có lịch sử)"

    user_content = (
        f"Lịch sử hội thoại:\n{conversation_history}\n\n"
        f"Câu hỏi hiện tại: {query}"
    )
    if formatted_docs:
        user_content += f"\n\nTài liệu tham khảo:\n{formatted_docs}"

    response = non_reasoning_llm.invoke([
        SystemMessage(content=RESPONSE_SYSTEM_PROMPT.format(current_time=_get_current_time())),
        HumanMessage(content=user_content),
    ])
    return {"messages": [response]}


def _route_after_router(state: State) -> Literal["rule_response", "orchestrator", "response"]:
    if state.get("response_type") in {"greeting", "out_of_scope"}:
        return "rule_response"
    return "orchestrator" if state.get("needs_data") else "response"


def _route_after_orchestrator(state: State) -> Literal["tool_node", "response"]:
    last_message = state["messages"][-1]
    has_tool_calls = bool(getattr(last_message, "tool_calls", None))
    under_limit = (state.get("iter_count") or 0) < MAX_ITERS
    return "tool_node" if (has_tool_calls and under_limit) else "response"


builder = StateGraph(State)
builder.add_node("router", router_node)
builder.add_node("orchestrator", orchestrator_node)
builder.add_node("tool_node", tool_node)
builder.add_node("response", response_node)
builder.add_node("rule_response", rule_response_node)

builder.set_entry_point("router")
builder.add_conditional_edges("router", _route_after_router)
builder.add_conditional_edges("orchestrator", _route_after_orchestrator)
builder.add_edge("tool_node", "orchestrator")
builder.add_edge("response", END)
builder.add_edge("rule_response", END)

graph = builder.compile()