app/nodes/router.py

"""
app/nodes/router.py — CHECKPOINT 3: CONDITIONAL ROUTING
LLM-based semantic router that classifies query into rag / tool / general.
"""

import json
from langchain_core.messages import HumanMessage
from app.state import AgentState
from app.tools import ALL_TOOLS
from app.utils.llm import llm


def router_node(state: AgentState) -> AgentState:
    tool_descriptions = "\n".join(
        f'- "{t.name}": {t.description}' for t in ALL_TOOLS
    )
    router_prompt = f"""You are a query router. Your ONLY job is to pick the correct route.

    ROUTE RULES — follow strictly in this order:

    1. "tool" — the query needs LIVE or REAL-TIME data that only a tool can provide.
    Use this if the query mentions: weather, temperature, forecast, calculate, math, arithmetic.
    Available tools:
    {tool_descriptions}

    2. "rag" — the query asks about AI/ML concepts that exist in the knowledge base.
    RAG topics: LangGraph, LangChain, RAG, Guardrails, HITL, Memory, Conditional routing.
    Do NOT use rag for weather or calculations — those always go to tool.

    3. "general" — everything else (greetings, opinions, open-ended questions).

    EXAMPLES:
    "What is the weather in Mumbai?" → tool (needs live weather data)
    "Calculate 25 * 48"             → tool (needs calculator)
    "What is LangGraph?"            → rag  (AI concept in knowledge base)
    "Tell me a joke"                → general

    Respond ONLY with valid JSON, no markdown:
    {{"route": "<tool|rag|general>", "reason": "<one sentence why>"}}

    User query: {state["query"]}"""

    try:
        response = llm.invoke([HumanMessage(content=router_prompt)])
        raw = response.content.strip().removeprefix("```json").removesuffix("```").strip()
        parsed = json.loads(raw)
        route = parsed.get("route", "general")
        reason = parsed.get("reason", "")
        if route not in ("rag", "tool", "general"):
            route = "general"
        print(f"[ROUTER] → '{route}' | {reason}")
    except Exception as e:
        print(f"[ROUTER] Failed ({e}), defaulting to 'general'.")
        route = "general"

    log = state.get("node_log", []) + [f"router → {route}"]
    return {**state, "route": route, "node_log": log}


def route_selector(state: AgentState) -> str:
    return state["route"]