graph_main.py

from agents.rate_agent import rate_manager
from agents.inventory_agent import inventory_manager
from agents.reservation_agent import reservation_manager
from agents.unknown_agent import unknown_handler
from agents.llm import llm
from langgraph.graph import StateGraph,START,END
from typing import TypedDict,Optional
from langgraph.checkpoint.memory import MemorySaver
from dotenv import load_dotenv
load_dotenv()



class SupportState(TypedDict):
    query: str
    next_agent: str
    response: str
    pending_manager: Optional[str]
    pending_action: Optional[dict]
    phone_number_id: str
    booking_link: str
    booking_step: str
    conversation_language: Optional[str]

    
memory = MemorySaver()


async def choose_agent(state: SupportState):
    if state.get("pending_manager"):
        return {"next_agent":state["pending_manager"]}
    
    prompt = f"""Classify the hotel owner's request into one of these 3 managers:
    'inventory' - Check availability, block, unblock, release, or unrelease rooms.
    'rates' - Check room prices or change room rates.
    'reservations' - Make a new booking, modify existing, or cancel.
    'unknown' - Any request not related to hotel management (coding, general knowledge, etc.)    
    Request: {state['query']}

    IMPORTANT RULES:
    - If the user is asking "how many" (कितने) or asking for counts, it is ALWAYS "inventory", even if they use the word "booked".
    - ONLY return "reservations" if there is a specific guest involved or an intent to create a new booking right now.
    Return ONLY one word: inventory, rates, reservations, or unknown."""

    response = llm.invoke(prompt)
    manager = response.content.strip().lower().replace(".","").replace("'","")

    if hasattr(response, 'usage_metadata') and response.usage_metadata:
        print(f"Tokens used: {response.usage_metadata}")

    if manager not in ["inventory","rates","reservations"]:
        manager = "unknown"
    print(f"LLM returned: '{response.content.strip().lower()}'")

    return {"next_agent": manager}


def agent_router(state: SupportState):
    return state["next_agent"]



graph = StateGraph(SupportState)
graph.add_node("choose_agent",choose_agent)
graph.add_node("inventory_manager",inventory_manager)
graph.add_node("rate_manager",rate_manager)
graph.add_node("reservation_manager",reservation_manager)
graph.add_node("unknown_handler",unknown_handler)


graph.add_edge(START,"choose_agent")
graph.add_conditional_edges(
    "choose_agent",
    agent_router,
    {
        "inventory": "inventory_manager",
        "rates": "rate_manager",
        "reservations": "reservation_manager",
        "unknown": "unknown_handler"
    }
)

graph.add_edge("inventory_manager", END)
graph.add_edge("rate_manager", END)
graph.add_edge("reservation_manager", END)
graph.add_edge("unknown_handler",END)

graph = graph.compile(checkpointer=memory)