from typing import TypedDict
from langgraph.graph import StateGraph, END
class State(TypedDict):
    query: str
    intent: str
    response: str
def classify(state):
    q = state["query"].lower()

    if "admission" in q:
        state["intent"] = "admission"
    elif "exam" in q:
        state["intent"] = "exam"
    elif "fee" in q or "fees" in q:
        state["intent"] = "fees"
    else:
        state["intent"] = "scholarship"

    return state
def admission(state):
    state["response"] = "Admission starts in July. Apply online with required documents."
    return state

def exam(state):
    state["response"] = "Entrance exam will be held in August."
    return state

def fees(state):
    state["response"] = "B.Tech yearly fee is ₹80,000."
    return state

def scholarship(state):
    state["response"] = "Scholarships are available based on merit and income."
    return state
def route(state):
    return state["intent"]
builder = StateGraph(State)

builder.add_node("Classifier", classify)
builder.add_node("Admission", admission)
builder.add_node("Exam", exam)
builder.add_node("Fees", fees)
builder.add_node("Scholarship", scholarship)

builder.set_entry_point("Classifier")

builder.add_conditional_edges(
    "Classifier",
    route,
    {
        "admission": "Admission",
        "exam": "Exam",
        "fees": "Fees",
        "scholarship": "Scholarship",
    },
)

builder.add_edge("Admission", END)
builder.add_edge("Exam", END)
builder.add_edge("Fees", END)
builder.add_edge("Scholarship", END)

graph = builder.compile()
query = input("Ask your question: ")

result = graph.invoke(
    {
        "query": query,
        "intent": "",
        "response": "",
    }
)

print("\nResponse:", result["response"])