elvismu/EmailClassificationLangGraph

1import gradio as gr2from typing import TypedDict, List, Dict, Any, Optional3from langgraph.graph import StateGraph, START, END4from langchain_openai import ChatOpenAI5from langchain_core.messages import HumanMessage6from langgraph.graph.state import CompiledStateGraph7 8model = ChatOpenAI(model="gpt-4o", temperature=0)9 10 11class EmailState(TypedDict):12    email: Dict[str, Any]13    is_spam: Optional[bool]14    draft_response: Optional[str]15    messages: List[Dict[str, Any]]16 17 18# Define nodes19def read_email(state: EmailState):20    email = state["email"]21    print(f"Alfred is processing an email from {email['sender']} with subject: {email['subject']}")22    return {}23 24 25def classify_email(state: EmailState):26    email = state["email"]27 28    prompt = f"""29As Alfred the butler of Mr wayne and it's SECRET identity Batman, analyze this email and determine if it is spam or legitimate and should be brought to Mr wayne's attention.30 31Email:32From: {email['sender']}33Subject: {email['subject']}34Body: {email['body']}35 36First, determine if this email is spam.37answer with SPAM or HAM if it's legitimate. Only reurn the answer38Answer :39    """40    messages = [HumanMessage(content=prompt)]41    response = model.invoke(messages)42 43    response_text = response.content.lower()44    print(response_text)45    is_spam = "spam" in response_text and "ham" not in response_text46 47    if not is_spam:48        new_messages = state.get("messages", []) + [49            {"role": "user", "content": prompt},50            {"role": "assistant", "content": response.content}51        ]52    else:53        new_messages = state.get("messages", [])54 55    return {56        "is_spam": is_spam,57        "messages": new_messages58    }59 60 61def handle_spam(state: EmailState):62    print(f"Alfred has marked the email as spam.")63    print("The email has been moved to the spam folder.")64    return {}65 66 67def drafting_response(state: EmailState):68    email = state["email"]69 70    prompt = f"""71As Alfred the butler, draft a polite preliminary response to this email.72 73Email:74From: {email['sender']}75Subject: {email['subject']}76Body: {email['body']}77 78Draft a brief, professional response that Mr. Wayne can review and personalize before sending.79    """80 81    messages = [HumanMessage(content=prompt)]82    response = model.invoke(messages)83 84    new_messages = state.get("messages", []) + [85        {"role": "user", "content": prompt},86        {"role": "assistant", "content": response.content}87    ]88 89    return {90        "draft_response": response.content,91        "messages": new_messages92    }93 94 95def notify_mr_wayne(state: EmailState):96    email = state["email"]97 98    print("\n" + "=" * 50)99    print(f"Sir, you've received an email from {email['sender']}.")100    print(f"Subject: {email['subject']}")101    print("\nI've prepared a draft response for your review:")102    print("-" * 50)103    print(state["draft_response"])104    print("=" * 50 + "\n")105 106    return {}107 108 109# Define routing logic110def route_email(state: EmailState) -> str:111    if state["is_spam"]:112        return "spam"113    else:114        return "legitimate"115 116 117# Create the graph118email_graph = StateGraph(EmailState)119 120# Add nodes121email_graph.add_node("read_email", read_email)  # the read_email node executes the read_mail function122email_graph.add_node("classify_email", classify_email)  # the classify_email node will execute the classify_email function123email_graph.add_node("handle_spam", handle_spam)  # same logic124email_graph.add_node("drafting_response", drafting_response)  # same logic125email_graph.add_node("notify_mr_wayne", notify_mr_wayne)  # same logic126 127# Add edges128email_graph.add_edge(START, "read_email") # After starting we go to the "read_email" node129 130email_graph.add_edge("read_email", "classify_email") # after_reading we classify131 132# Add conditional edges133email_graph.add_conditional_edges(134    "classify_email", # after classify, we run the "route_email" function"135    route_email,136    {137        "spam": "handle_spam", # if it returns "Spam", we go the "handle_span" node138        "legitimate": "drafting_response" # and if it's legitimate, we go to the "drafting response" node139    }140)141 142# Add final edges143email_graph.add_edge("handle_spam", END) # after handling spam we always end144email_graph.add_edge("drafting_response", "notify_mr_wayne")145email_graph.add_edge("notify_mr_wayne", END) # after notifying Me wayne, we can end  too146 147compiled_graph = email_graph.compile()148 149def classify(email: str) -> CompiledStateGraph:150    compiled_graph.invoke({151        "email": email,152        "is_spam": None,153        "draft_response": None,154        "messages": []155    })156    return compiled_graph157 158demo = gr.ChatInterface(159    classify,160    type="messages",161    title="EmailClassifyingLangGraph",162    description="Classify your emails!",163    multimodal=True)164 165demo.launch()166