Multi-Agent Bank Support (Python + LangGraph)

import os
import warnings

# 1. SUPPRESS WARNINGS (Must be done before other imports)
warnings.filterwarnings("ignore")
os.environ["PYTHONWARNINGS"] = "ignore"
os.environ["GRPC_VERBOSITY"] = "NONE"

from typing import TypedDict
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_react_agent

# 2. LOAD ENVIRONMENT VARIABLES
load_dotenv()
api_key = os.getenv("GOOGLE_API_KEY")

if not api_key:
    print("ERROR: GOOGLE_API_KEY not found in .env file.")
    exit()

llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0)

class AgentState(TypedDict):
    user_input: str
    category: str
    response: str

# 4. TOOLS DEFINITION
@tool
def check_account_balance(account_number: str) -> str:
    """Useful to check the balance of a bank customer using their Account Number."""
    if account_number == "123456789":
        return "The account balance is $1,500.00 and there are no late fees."
    return "Account not found or balance unavailable."

# 5. NODE DEFINITIONS

def router_node(state: AgentState) -> AgentState:
    prompt = ChatPromptTemplate.from_messages([
        ("system", """You are the Supervisor Agent for a Smart Bank.
        Analyze the user query and classify it into EXACTLY ONE of these categories:
        - 'technical': app crashes, login issues, website bugs.
        - 'billing': questions about fees, balances, transactions, and account checking.
        - 'general': business hours, general policies, branch locations.
        - 'escalate': angry customers, complex fraud, explicit requests for a human.
        - 'irrelevant': anything not related to banking or this bank.

        CRITICAL: Respond with ONLY the exact category word. No punctuation, no markdown."""),
        ("user", "{input}")
    ])
    chain = prompt | llm
    result = chain.invoke({"input": state["user_input"]})
    category = result.content.strip().lower().replace("*", "")
    return {"category": category}

def technical_agent_node(state: AgentState) -> AgentState:
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a Technical Support Agent. Guide the user to solve app/website issues politely."),
        ("user", "{input}")
    ])
    chain = prompt | llm
    result = chain.invoke({"input": state["user_input"]})
    return {"response": result.content}

def general_agent_node(state: AgentState) -> AgentState:
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a General Support Agent. Answer general banking questions concisely."),
        ("user", "{input}")
    ])
    chain = prompt | llm
    result = chain.invoke({"input": state["user_input"]})
    return {"response": result.content}

def billing_agent_node(state: AgentState) -> AgentState:
    billing_agent = create_react_agent(llm, tools=[check_account_balance])
    system_msg = "You are a Billing Agent. If the user provides an Account Number, ALWAYS use the check_account_balance tool."
    messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": state["user_input"]},
    ]
    result = billing_agent.invoke({"messages": messages})
    final_response = result["messages"][-1].content
    return {"response": final_response}

def escalate_node(state: AgentState) -> AgentState:
    return {"response": "I am transferring you to a human agent right now. Please hold."}
    
def irrelevant_node(state: AgentState) -> AgentState:
    return {"response": "I can only assist with bank-related queries. How can I help you with your finances today?"}

def route_decision(state: AgentState) -> str:
    category = state.get("category", "irrelevant")
    valid_categories = ["technical", "billing", "general", "escalate", "irrelevant"]
    return category if category in valid_categories else "irrelevant"

# 6. GRAPH CONSTRUCTION

workflow = StateGraph(AgentState)

workflow.add_node("router", router_node)
workflow.add_node("technical", technical_agent_node)
workflow.add_node("billing", billing_agent_node)
workflow.add_node("general", general_agent_node)
workflow.add_node("escalate", escalate_node)
workflow.add_node("irrelevant", irrelevant_node)

workflow.set_entry_point("router")

workflow.add_conditional_edges(
    "router",
    route_decision,
    {
        "technical": "technical",
        "billing": "billing",
        "general": "general",
        "escalate": "escalate",
        "irrelevant": "irrelevant",
    }
)

workflow.add_edge("technical", END)
workflow.add_edge("billing", END)
workflow.add_edge("general", END)
workflow.add_edge("escalate", END)
workflow.add_edge("irrelevant", END)

app = workflow.compile()

# 7. INTERACTIVE MAIN LOOP

def main():
    print("\n" + "="*60)
    print("SMART BANK AI SUPPORT SYSTEM")
    print("Type 'exit' to quit.")
    print("="*60)

    while True:
        user_query = input("\nYou: ")
        
        if user_query.lower() in ["exit", "quit"]:
            print("Goodbye!")
            break
        
        if not user_query.strip():
            continue

        try:
            result = app.invoke({"user_input": user_query})
            print(f"\n[SYSTEM] Category: {result.get('category', '').upper()}")
            print(f"[AGENT]  Response: {result.get('response')}")
            print("-" * 60)
        except Exception as e:
            print(f"\n[ERROR] Something went wrong: {e}")
            print("Please ensure your LangChain packages are up to date.")

if __name__ == "__main__":
    main()